diff --git a/Data/BitUtil.hs b/Data/BitUtil.hs new file mode 100644 --- /dev/null +++ b/Data/BitUtil.hs @@ -0,0 +1,71 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.BitUtil +-- Copyright : (c) Clark Gaebel 2012 +-- (c) Johan Tibel 2012 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +----------------------------------------------------------------------------- + +module Data.BitUtil + ( highestBitMask + , shiftLL + , shiftRL + ) where + +-- On GHC, include MachDeps.h to get WORD_SIZE_IN_BITS macro. +#if defined(__GLASGOW_HASKELL__) +# include "MachDeps.h" +#endif + +import Data.Bits ((.|.), xor) + +#if __GLASGOW_HASKELL__ +import GHC.Exts (Word(..), Int(..)) +import GHC.Prim (uncheckedShiftL#, uncheckedShiftRL#) +#else +import Data.Word (shiftL, shiftR) +#endif + +-- The highestBitMask implementation is based on +-- http://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2 +-- which has been put in the public domain. + +-- | Return a word where only the highest bit is set. +highestBitMask :: Word -> Word +highestBitMask x1 = let x2 = x1 .|. x1 `shiftRL` 1 + x3 = x2 .|. x2 `shiftRL` 2 + x4 = x3 .|. x3 `shiftRL` 4 + x5 = x4 .|. x4 `shiftRL` 8 + x6 = x5 .|. x5 `shiftRL` 16 +#if !(defined(__GLASGOW_HASKELL__) && WORD_SIZE_IN_BITS==32) + x7 = x6 .|. x6 `shiftRL` 32 + in x7 `xor` (x7 `shiftRL` 1) +#else + in x6 `xor` (x6 `shiftRL` 1) +#endif +{-# INLINE highestBitMask #-} + +-- Right and left logical shifts. +shiftRL, shiftLL :: Word -> Int -> Word +#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 #-} diff --git a/Data/Graph.hs b/Data/Graph.hs new file mode 100644 --- /dev/null +++ b/Data/Graph.hs @@ -0,0 +1,478 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE Rank2Types #-} +#endif +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- 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 +-- +-- A version of the graph algorithms described in: +-- +-- /Structuring Depth-First Search Algorithms in Haskell/, +-- by David King and John Launchbury. +-- +----------------------------------------------------------------------------- + +module Data.Graph( + + -- * External interface + + -- At present the only one with a "nice" external interface + stronglyConnComp, stronglyConnCompR, SCC(..), flattenSCC, flattenSCCs, + + -- * Graphs + + Graph, Table, Bounds, Edge, Vertex, + + -- ** Building graphs + + graphFromEdges, graphFromEdges', buildG, transposeG, + -- reverseE, + + -- ** Graph properties + + vertices, edges, + outdegree, indegree, + + -- * Algorithms + + dfs, dff, + topSort, + components, + scc, + bcc, + -- tree, back, cross, forward, + reachable, path, + + module Data.Tree + + ) where + +#if __GLASGOW_HASKELL__ +# define USE_ST_MONAD 1 +#endif + +-- Extensions +#if USE_ST_MONAD +import Control.Monad.ST +import Data.Array.ST (STArray, newArray, readArray, writeArray) +#else +import Data.IntSet (IntSet) +import qualified Data.IntSet as Set +#endif +import Data.Tree (Tree(Node), Forest) + +-- std interfaces +import Control.Applicative +import Control.DeepSeq (NFData(rnf)) +import Data.Maybe +import Data.Array +import Data.List + +------------------------------------------------------------------------- +-- - +-- 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. + +instance NFData a => NFData (SCC a) where + rnf (AcyclicSCC v) = rnf v + rnf (CyclicSCC vs) = rnf vs + +instance Functor SCC where + fmap f (AcyclicSCC v) = AcyclicSCC (f v) + fmap f (CyclicSCC vs) = CyclicSCC (fmap f vs) + +-- | The vertices of a list of strongly connected components. +flattenSCCs :: [SCC a] -> [a] +flattenSCCs = concatMap flattenSCC + +-- | The vertices of a strongly connected component. +flattenSCC :: SCC vertex -> [vertex] +flattenSCC (AcyclicSCC v) = [v] +flattenSCC (CyclicSCC vs) = vs + +-- | 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] + +stronglyConnComp edges0 + = map get_node (stronglyConnCompR edges0) + where + get_node (AcyclicSCC (n, _, _)) = AcyclicSCC n + get_node (CyclicSCC triples) = CyclicSCC [n | (n,_,_) <- triples] + +-- | The strongly connected components of a directed graph, topologically +-- sorted. The function is the same as 'stronglyConnComp', except that +-- all the information about each node retained. +-- This interface is used when you expect to apply 'SCC' to +-- (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 + +stronglyConnCompR [] = [] -- added to avoid creating empty array in graphFromEdges -- SOF +stronglyConnCompR edges0 + = map decode forest + where + (graph, vertex_fn,_) = graphFromEdges edges0 + forest = scc graph + decode (Node v []) | mentions_itself v = CyclicSCC [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 + mentions_itself v = v `elem` (graph ! v) + +------------------------------------------------------------------------- +-- - +-- Graphs +-- - +------------------------------------------------------------------------- + +-- | Abstract representation of vertices. +type Vertex = Int +-- | Table indexed by a contiguous set of vertices. +type Table a = Array Vertex a +-- | Adjacency list representation of a graph, mapping each vertex to its +-- list of successors. +type Graph = Table [Vertex] +-- | The bounds of a 'Table'. +type Bounds = (Vertex, Vertex) +-- | An edge from the first vertex to the second. +type Edge = (Vertex, Vertex) + +-- | All vertices of a graph. +vertices :: Graph -> [Vertex] +vertices = indices + +-- | All edges of a graph. +edges :: Graph -> [Edge] +edges g = [ (v, w) | v <- vertices g, w <- g!v ] + +mapT :: (Vertex -> a -> b) -> Table a -> Table b +mapT f t = array (bounds t) [ (,) v (f v (t!v)) | v <- indices t ] + +-- | Build a graph from a list of edges. +buildG :: Bounds -> [Edge] -> Graph +buildG bounds0 edges0 = accumArray (flip (:)) [] bounds0 edges0 + +-- | The graph obtained by reversing all edges. +transposeG :: Graph -> Graph +transposeG g = buildG (bounds g) (reverseE g) + +reverseE :: Graph -> [Edge] +reverseE g = [ (w, v) | (v, w) <- edges g ] + +-- | A table of the count of edges from each node. +outdegree :: Graph -> Table Int +outdegree = mapT numEdges + where numEdges _ ws = length ws + +-- | A table of the count of edges into each node. +indegree :: Graph -> Table Int +indegree = outdegree . transposeG + +-- | Identical to 'graphFromEdges', except that the return value +-- 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])) +graphFromEdges' x = (a,b) where + (a,b,_) = graphFromEdges x + +-- | Build a graph from a list of nodes uniquely identified by keys, +-- with a list of keys of nodes this node should have edges to. +-- 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) +graphFromEdges edges0 + = (graph, \v -> vertex_map ! v, key_vertex) + where + max_v = length edges0 - 1 + bounds0 = (0,max_v) :: (Vertex, Vertex) + sorted_edges = sortBy lt edges0 + 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 + + (_,k1,_) `lt` (_,k2,_) = k1 `compare` k2 + + -- key_vertex :: key -> Maybe Vertex + -- 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 + +------------------------------------------------------------------------- +-- - +-- Depth first search +-- - +------------------------------------------------------------------------- + +-- | A spanning forest of the graph, obtained from a depth-first search of +-- the graph starting from each vertex in an unspecified order. +dff :: Graph -> Forest Vertex +dff g = dfs g (vertices g) + +-- | A spanning forest of the part of the graph reachable from the listed +-- vertices, obtained from a depth-first search of the graph starting at +-- each of the listed vertices in order. +dfs :: Graph -> [Vertex] -> Forest Vertex +dfs g vs = prune (bounds g) (map (generate g) vs) + +generate :: Graph -> Vertex -> Tree Vertex +generate g v = Node v (map (generate g) (g!v)) + +prune :: Bounds -> Forest Vertex -> Forest Vertex +prune bnds ts = run bnds (chop ts) + +chop :: Forest Vertex -> SetM s (Forest Vertex) +chop [] = return [] +chop (Node v ts : us) + = do + visited <- contains v + if visited then + chop us + else do + include v + as <- chop ts + bs <- chop us + return (Node v as : bs) + +-- A monad holding a set of vertices visited so far. +#if USE_ST_MONAD + +-- Use the ST monad if available, for constant-time primitives. + +newtype SetM s a = SetM { runSetM :: STArray s Vertex Bool -> ST s a } + +instance Monad (SetM s) where + return x = SetM $ const (return x) + {-# INLINE return #-} + SetM v >>= f = SetM $ \s -> do { x <- v s; runSetM (f x) s } + {-# INLINE (>>=) #-} + +instance Functor (SetM s) where + f `fmap` SetM v = SetM $ \s -> f `fmap` v s + {-# INLINE fmap #-} + +instance Applicative (SetM s) where + pure x = SetM $ const (return x) + {-# INLINE pure #-} + SetM f <*> SetM v = SetM $ \s -> f s >>= (`fmap` v s) + -- We could also use the following definition + -- SetM f <*> SetM v = SetM $ \s -> f s <*> v s + -- but Applicative (ST s) instance is present only in GHC 7.2+ + {-# INLINE (<*>) #-} + +run :: Bounds -> (forall s. SetM s a) -> a +run bnds act = runST (newArray bnds False >>= runSetM act) + +contains :: Vertex -> SetM s Bool +contains v = SetM $ \ m -> readArray m v + +include :: Vertex -> SetM s () +include v = SetM $ \ m -> writeArray m v True + +#else /* !USE_ST_MONAD */ + +-- Portable implementation using IntSet. + +newtype SetM s a = SetM { runSetM :: IntSet -> (a, IntSet) } + +instance Monad (SetM s) where + return x = SetM $ \s -> (x, s) + SetM v >>= f = SetM $ \s -> case v s of (x, s') -> runSetM (f x) s' + +instance Functor (SetM s) where + f `fmap` SetM v = SetM $ \s -> case v s of (x, s') -> (f x, s') + {-# INLINE fmap #-} + +instance Applicative (SetM s) where + pure x = SetM $ \s -> (x, s) + {-# INLINE pure #-} + SetM f <*> SetM v = SetM $ \s -> case f s of (k, s') -> case v s' of (x, s'') -> (k x, s'') + {-# INLINE (<*>) #-} + +run :: Bounds -> SetM s a -> a +run _ act = fst (runSetM act Set.empty) + +contains :: Vertex -> SetM s Bool +contains v = SetM $ \ m -> (Set.member v m, m) + +include :: Vertex -> SetM s () +include v = SetM $ \ m -> ((), Set.insert v m) + +#endif /* !USE_ST_MONAD */ + +------------------------------------------------------------------------- +-- - +-- Algorithms +-- - +------------------------------------------------------------------------- + +------------------------------------------------------------ +-- Algorithm 1: depth first search numbering +------------------------------------------------------------ + +preorder' :: Tree a -> [a] -> [a] +preorder' (Node a ts) = (a :) . preorderF' ts + +preorderF' :: Forest a -> [a] -> [a] +preorderF' ts = foldr (.) id $ map preorder' ts + +preorderF :: Forest a -> [a] +preorderF ts = preorderF' ts [] + +tabulate :: Bounds -> [Vertex] -> Table Int +tabulate bnds vs = array bnds (zipWith (,) vs [1..]) + +preArr :: Bounds -> Forest Vertex -> Table Int +preArr bnds = tabulate bnds . preorderF + +------------------------------------------------------------ +-- Algorithm 2: topological sorting +------------------------------------------------------------ + +postorder :: Tree a -> [a] -> [a] +postorder (Node a ts) = postorderF ts . (a :) + +postorderF :: Forest a -> [a] -> [a] +postorderF ts = foldr (.) id $ map postorder ts + +postOrd :: Graph -> [Vertex] +postOrd g = postorderF (dff g) [] + +-- | A topological sort of the graph. +-- The order is partially specified by the condition that a vertex /i/ +-- precedes /j/ whenever /j/ is reachable from /i/ but not vice versa. +topSort :: Graph -> [Vertex] +topSort = reverse . postOrd + +------------------------------------------------------------ +-- Algorithm 3: connected components +------------------------------------------------------------ + +-- | The connected components of a graph. +-- Two vertices are connected if there is a path between them, traversing +-- edges in either direction. +components :: Graph -> Forest Vertex +components = dff . undirected + +undirected :: Graph -> Graph +undirected g = buildG (bounds g) (edges g ++ reverseE g) + +-- Algorithm 4: strongly connected components + +-- | The strongly connected components of a graph. +scc :: Graph -> Forest Vertex +scc g = dfs g (reverse (postOrd (transposeG g))) + +------------------------------------------------------------ +-- Algorithm 5: Classifying edges +------------------------------------------------------------ + +{- +XXX unused code + +tree :: Bounds -> Forest Vertex -> Graph +tree bnds ts = buildG bnds (concat (map flat ts)) + where flat (Node v ts') = [ (v, w) | Node w _us <- ts' ] + ++ concat (map flat ts') + +back :: Graph -> Table Int -> Graph +back g post = mapT select g + where select v ws = [ w | w <- ws, post!v < post!w ] + +cross :: Graph -> Table Int -> Table Int -> Graph +cross g pre post = mapT select g + where select v ws = [ w | w <- ws, post!v > post!w, pre!v > pre!w ] + +forward :: Graph -> Graph -> Table Int -> Graph +forward g tree' pre = mapT select g + where select v ws = [ w | w <- ws, pre!v < pre!w ] \\ tree' ! v +-} + +------------------------------------------------------------ +-- Algorithm 6: Finding reachable vertices +------------------------------------------------------------ + +-- | A list of vertices reachable from a given vertex. +reachable :: Graph -> Vertex -> [Vertex] +reachable g v = preorderF (dfs g [v]) + +-- | Is the second vertex reachable from the first? +path :: Graph -> Vertex -> Vertex -> Bool +path g v w = w `elem` (reachable g v) + +------------------------------------------------------------ +-- Algorithm 7: Biconnected components +------------------------------------------------------------ + +-- | The biconnected components of a graph. +-- An undirected graph is biconnected if the deletion of any vertex +-- leaves it connected. +bcc :: Graph -> Forest [Vertex] +bcc g = (concat . map bicomps . map (do_label g dnum)) forest + where forest = dff g + dnum = preArr (bounds g) forest + +do_label :: Graph -> Table Int -> Tree Vertex -> Tree (Vertex,Int,Int) +do_label g dnum (Node v ts) = Node (v,dnum!v,lv) us + where us = map (do_label g dnum) ts + lv = minimum ([dnum!v] ++ [dnum!w | w <- g!v] + ++ [lu | Node (_,_,lu) _ <- us]) + +bicomps :: Tree (Vertex,Int,Int) -> Forest [Vertex] +bicomps (Node (v,_,_) ts) + = [ Node (v:vs) us | (_,Node vs us) <- map collect ts] + +collect :: Tree (Vertex,Int,Int) -> (Int, Tree [Vertex]) +collect (Node (v,dv,lv) ts) = (lv, Node (v:vs) cs) + where collected = map collect ts + vs = concat [ ws | (lw, Node ws _) <- collected, lw= 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 deprecated value strict functions. Please note that these functions +-- have different strictness properties than those in "Data.IntMap.Strict": +-- they only evaluate the result of the combining function. For example, the +-- default value to 'insertWith'' is only evaluated if the combining function +-- is called and uses it. +-- +-- 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 () -- hide foldr +import qualified Data.IntMap.Strict as Strict +import Data.IntMap.Lazy + +-- | /Deprecated./ As of version 0.5, replaced by +-- 'Data.IntMap.Strict.insertWith'. +-- +-- /O(log n)/. Same as 'insertWith', but the result of the combining function +-- is evaluated to WHNF before inserted to the map. + +insertWith' :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWith' = Strict.insertWith + +-- | /Deprecated./ As of version 0.5, replaced by +-- 'Data.IntMap.Strict.insertWithKey'. +-- +-- /O(log n)/. Same as 'insertWithKey', but the result of the combining +-- function is evaluated to WHNF before inserted to the map. + +insertWithKey' :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithKey' = Strict.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 :: (Key -> 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,2197 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +{-# LANGUAGE ScopedTypeVariables #-} +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE TypeFamilies #-} +#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 + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * 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 + , link + , bin + , zero + , nomatch + , match + , mask + , maskW + , shorter + , branchMask + , highestBitMask + , foldlStrict + ) where + +import Control.Applicative (Applicative(pure, (<*>)), (<$>)) +import Control.DeepSeq (NFData(rnf)) +import Control.Monad (liftM) +import Data.Bits +import qualified Data.Foldable as Foldable +import Data.Maybe (fromMaybe) +import Data.Monoid (Monoid(..)) +import Data.Traversable (Traversable(traverse)) +import Data.Typeable +import Data.Word (Word) +import Prelude hiding (lookup, map, filter, foldr, foldl, null) + +import Data.BitUtil +import Data.IntSet.Base (Key) +import qualified Data.IntSet.Base as IntSet +import Data.StrictPair + +#if __GLASGOW_HASKELL__ +import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix), + DataType, mkDataType) +import GHC.Exts (build) +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts as GHCExts +#endif +import Text.Read +#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 #-} + +{-------------------------------------------------------------------- + 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 + +{-------------------------------------------------------------------- + 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 t = go t + where go Nil = mempty + go (Tip _ v) = v + go (Bin _ _ l r) = go l `mappend` go r + {-# INLINABLE fold #-} + foldr = foldr + {-# INLINE foldr #-} + foldl = foldl + {-# INLINE foldl #-} + foldMap f t = go t + where go Nil = mempty + go (Tip _ v) = f v + go (Bin _ _ l r) = go l `mappend` go r + {-# INLINE foldMap #-} + +instance Traversable IntMap where + traverse f = traverseWithKey (\_ -> f) + {-# INLINE traverse #-} + +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 provide limited 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 _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = intMapDataType + dataCast1 f = gcast1 f + +fromListConstr :: Constr +fromListConstr = mkConstr intMapDataType "fromList" [] Prefix + +intMapDataType :: DataType +intMapDataType = mkDataType "Data.IntMap.Base.IntMap" [fromListConstr] + +#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 + +infix 4 `member` + +-- | /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 + +infix 4 `notMember` + +-- | /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 -> link k (Tip k x) p t + | zero k m -> Bin p m (insert k x l) r + | otherwise -> Bin p m l (insert k x r) + Tip ky _ + | k==ky -> Tip k x + | otherwise -> link k (Tip k x) ky t + Nil -> Tip k x + +-- 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 -> link k (Tip k x) p t + | zero k m -> Bin p m (insertWithKey f k x l) r + | otherwise -> Bin p m l (insertWithKey f k x r) + Tip ky y + | k==ky -> Tip k (f k x y) + | otherwise -> link k (Tip k x) ky t + Nil -> Tip k x + +-- | /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,link k (Tip k x) p t) + | zero k m -> let (found,l') = insertLookupWithKey f k x l in (found,Bin p m l' r) + | otherwise -> let (found,r') = insertLookupWithKey f k x r in (found,Bin p m l r') + Tip ky y + | k==ky -> (Just y,Tip k (f k x y)) + | otherwise -> (Nothing,link k (Tip k x) ky t) + Nil -> (Nothing,Tip k x) + + +{-------------------------------------------------------------------- + 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 -> link k (Tip k x) p t + | zero k m -> bin p m (alter f k l) r + | otherwise -> bin p m l (alter f k r) + Tip ky y + | k==ky -> case f (Just y) of + Just x -> Tip ky x + Nothing -> Nil + | otherwise -> case f Nothing of + Just x -> link k (Tip k x) ky t + Nothing -> Tip ky y + Nil -> case f Nothing of + Just x -> Tip k x + Nothing -> Nil + + +{-------------------------------------------------------------------- + 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 + +infixl 5 `union` + +-- | /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 + +infixl 5 `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 -> 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_link p1 (g1 t1) p2 (g2 t2) + where + merge1 | nomatch p2 p1 m1 = maybe_link 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_link 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_link 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_link 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_link 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_link k1 (g1 t1) k2 (g2 t2) + merge t1 _ Nil = g1 t1 + + go Nil t2 = g2 t2 + + maybe_link _ Nil _ t2 = t2 + maybe_link _ t1 _ Nil = t1 + maybe_link p1 t1 p2 t2 = link p1 t1 p2 t2 + {-# INLINE maybe_link #-} +{-# 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. Returns an empty map if the map is empty. +-- +-- Note that this is a change of behaviour for consistency with 'Data.Map.Map' – +-- versions prior to 0.5 threw an error if the 'IntMap' was already empty. +deleteMin :: IntMap a -> IntMap a +deleteMin = maybe Nil snd . minView + +-- | /O(min(n,W))/. Delete the maximal key. Returns an empty map if the map is empty. +-- +-- Note that this is a change of behaviour for consistency with 'Data.Map.Map' – +-- versions prior to 0.5 threw an error if the 'IntMap' was already empty. +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 +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 +{-# INLINE traverseWithKey #-} + +-- | /O(n)/. The function @'mapAccum'@ threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a b = (a ++ b, b ++ "X") +-- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) + +mapAccum :: (a -> b -> (a,c)) -> a -> 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 predicate0 t0 = toPair $ go predicate0 t0 + where + go predicate t + = case t of + Bin p m l r + -> let (l1 :*: l2) = go predicate l + (r1 :*: r2) = go 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 f0 t0 = toPair $ go f0 t0 + where + go f (Bin p m l r) + = bin p m l1 r1 :*: bin p m l2 r2 + where + (l1 :*: l2) = go f l + (r1 :*: r2) = go f r + go f (Tip k x) = case f k x of + Left y -> (Tip k y :*: Nil) + Right z -> (Nil :*: Tip k z) + go _ 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) -> let lt' = union r lt + in lt' `seq` (lt', gt) + else case go k r of (lt :*: gt) -> let gt' = union gt l + in gt' `seq` (lt, gt') + _ -> case go k t of + (lt :*: gt) -> (lt, gt) + 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) -> let lt' = union r lt + in lt' `seq` (lt', fnd, gt) + else case go k r of + (lt, fnd, gt) -> let gt' = union gt l + in gt' `seq` (lt, fnd, gt') + _ -> go k t + where + go k' t'@(Bin p m l r) + | nomatch k' p m = if k' > p then (t', Nothing, Nil) else (Nil, Nothing, t') + | zero k' m = case go k' l of + (lt, fnd, gt) -> let gt' = union gt r in gt' `seq` (lt, fnd, gt') + | otherwise = case go k' r of + (lt, fnd, gt) -> let lt' = union l lt in lt' `seq` (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 :: (Key -> 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' :: (Key -> 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 -> Key -> 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 -> Key -> 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' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given monoid, such that +-- +-- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@ +-- +-- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids. +foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m +foldMapWithKey f = go + where + go Nil = mempty + go (Tip kx x) = f kx x + go (Bin _ _ l r) = go l `mappend` go r +{-# INLINE foldMapWithKey #-} + +{-------------------------------------------------------------------- + List variations +--------------------------------------------------------------------} +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- Subject to list fusion. +-- +-- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] +-- > elems empty == [] + +elems :: 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 +--------------------------------------------------------------------} +#if __GLASGOW_HASKELL__ >= 708 +instance GHCExts.IsList (IntMap a) where + type Item (IntMap a) = (Key,a) + fromList = fromList + toList = toList +#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 :: 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 :: forall a. [(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 (link 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 +--------------------------------------------------------------------} +{-------------------------------------------------------------------- + Link +--------------------------------------------------------------------} +link :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a +link 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 link #-} + +{-------------------------------------------------------------------- + @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 #-} + +{-------------------------------------------------------------------- + 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 #-} + +-- | /O(1)/. Decompose a map into pieces based on the structure of the underlying +-- tree. This function is useful for consuming a map in parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that the +-- pieces returned will be in ascending order (all elements in the first submap +-- less than all elements in the second, and so on). +-- +-- Examples: +-- +-- > splitRoot (fromList (zip [1..6::Int] ['a'..])) == +-- > [fromList [(1,'a'),(2,'b'),(3,'c')],fromList [(4,'d'),(5,'e'),(6,'f')]] +-- +-- > splitRoot empty == [] +-- +-- Note that the current implementation does not return more than two submaps, +-- but you should not depend on this behaviour because it can change in the +-- future without notice. +splitRoot :: IntMap a -> [IntMap a] +splitRoot orig = + case orig of + Nil -> [] + x@(Tip _ _) -> [x] + Bin _ m l r | m < 0 -> [r, l] + | otherwise -> [l, r] +{-# INLINE splitRoot #-} + + +{-------------------------------------------------------------------- + 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,217 @@ +{-# 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 + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , IM.filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * 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,978 @@ +{-# 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 + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * 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 Data.BitUtil +import qualified Data.IntSet.Base as IntSet +import Data.StrictPair + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. +-- +-- Here's an example illustrating the first property: +-- +-- > delete undefined m == undefined +-- +-- Here are some examples that illustrate the second property: +-- +-- > map (\ v -> undefined) m == undefined -- m is not empty +-- > mapKeys (\ k -> undefined) m == undefined -- m is not empty + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} + +-- | /O(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 = 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 -> link k (Tip k x) p t + | zero k m -> Bin p m (insert k x l) r + | otherwise -> Bin p m l (insert k x r) + Tip ky _ + | k==ky -> Tip k x + | otherwise -> link k (Tip k x) ky t + Nil -> Tip k x + +-- 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` + case t of + Bin p m l r + | nomatch k p m -> link k (singleton k x) p t + | zero k m -> Bin p m (insertWithKey f k x l) r + | otherwise -> Bin p m l (insertWithKey f k x r) + Tip ky y + | k==ky -> Tip k $! f k x y + | otherwise -> link k (singleton k x) ky t + Nil -> singleton 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 f0 k0 x0 t0 = k0 `seq` toPair $ go f0 k0 x0 t0 + where + go f k x t = + case t of + Bin p m l r + | nomatch k p m -> Nothing :*: link k (singleton k x) p t + | zero k m -> let (found :*: l') = go f k x l in (found :*: Bin p m l' r) + | otherwise -> let (found :*: r') = go f k x r in (found :*: Bin p m l r') + Tip ky y + | k==ky -> (Just y :*: (Tip k $! f k x y)) + | otherwise -> (Nothing :*: link k (singleton k x) ky t) + Nil -> Nothing :*: (singleton 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 f0 k0 t0 = k0 `seq` toPair $ go f0 k0 t0 + where + go f k t = + case t of + Bin p m l r + | nomatch k p m -> (Nothing :*: t) + | zero k m -> let (found :*: l') = go f k l in (found :*: bin p m l' r) + | otherwise -> let (found :*: r') = go f k r in (found :*: bin p m l r') + Tip ky y + | k==ky -> case f k y of + Just y' -> y' `seq` (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 -> x `seq` link k (Tip k x) p t + | zero k m -> bin p m (alter f k l) r + | otherwise -> bin p m l (alter f k r) + 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` link 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 f0 a0 t0 = toPair $ go f0 a0 t0 + where + go f a t + = case t of + Bin p m l r -> let (a1 :*: l') = go f a l + (a2 :*: r') = go f a1 r + in (a2 :*: Bin p m l' r') + Tip k x -> let (a',x') = f a k x in x' `seq` (a' :*: Tip k x') + 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 f0 a0 t0 = toPair $ go f0 a0 t0 + where + go f a t + = case t of + Bin p m l r -> let (a1 :*: r') = go f a r + (a2 :*: l') = go f a1 l + in (a2 :*: Bin p m l' r') + Tip k x -> let (a',x') = f a k x in x' `seq` (a' :*: Tip k x') + Nil -> (a :*: 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 f0 t0 = toPair $ go f0 t0 + where + go f (Bin p m l r) + = bin p m l1 r1 :*: bin p m l2 r2 + where + (l1 :*: l2) = go f l + (r1 :*: r2) = go f r + go f (Tip k x) = case f k x of + Left y -> y `seq` (Tip k y :*: Nil) + Right z -> z `seq` (Nil :*: Tip k z) + go _ 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 (link 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 new file mode 100644 --- /dev/null +++ b/Data/IntSet.hs @@ -0,0 +1,150 @@ +{-# 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 + , Key + + -- * 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 + , splitRoot + + -- * 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,1534 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash, BangPatterns, DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE TypeFamilies #-} +#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(..), Key -- 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 + , splitRoot + + -- * 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 + +-- We want to be able to compile without cabal. Nevertheless +-- #if defined(MIN_VERSION_base) && MIN_VERSION_base(4,5,0) +-- does not work, because if MIN_VERSION_base is undefined, +-- the last condition is syntactically wrong. +#define MIN_VERSION_base_4_5_0 0 +#ifdef MIN_VERSION_base +#if MIN_VERSION_base(4,5,0) +#undef MIN_VERSION_base_4_5_0 +#define MIN_VERSION_base_4_5_0 1 +#endif +#endif + +#define MIN_VERSION_base_4_7_0 0 +#ifdef MIN_VERSION_base +#if MIN_VERSION_base(4,7,0) +#undef MIN_VERSION_base_4_7_0 +#define MIN_VERSION_base_4_7_0 1 +#endif +#endif + +import Control.DeepSeq (NFData) +import Data.Bits +import qualified Data.List as List +import Data.Maybe (fromMaybe) +import Data.Monoid (Monoid(..)) +import Data.Typeable +import Data.Word (Word) +import Prelude hiding (filter, foldr, foldl, null, map) + +import Data.BitUtil +import Data.StrictPair + +#if __GLASGOW_HASKELL__ +import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix), DataType, mkDataType) +import Text.Read +#endif + +#if __GLASGOW_HASKELL__ +import GHC.Exts (Int(..), build) +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts as GHCExts +#endif +import GHC.Prim (indexInt8OffAddr#) +#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 #-} + +{-------------------------------------------------------------------- + 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 +type Key = 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 provide limited reflection services for the sake of data abstraction. + +instance Data IntSet where + gfoldl f z is = z fromList `f` (toList is) + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = intSetDataType + +fromListConstr :: Constr +fromListConstr = mkConstr intSetDataType "fromList" [] Prefix + +intSetDataType :: DataType +intSetDataType = mkDataType "Data.IntSet.Base.IntSet" [fromListConstr] + +#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 :: Key -> 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 + +infix 4 `member` + +-- | /O(min(n,W))/. Is the element not in the set? +notMember :: Key -> IntSet -> Bool +notMember k = not . member k + +infix 4 `notMember` + +-- | /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 :: Key -> IntSet -> Maybe Key +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 :: Key -> IntSet -> Maybe Key +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 :: Key -> IntSet -> Maybe Key +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 :: Key -> IntSet -> Maybe Key +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 Key +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 Key +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 :: Key -> 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 :: Key -> 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 -> link kx (Tip kx bm) p t + | zero kx m -> Bin p m (insertBM kx bm l) r + | otherwise -> Bin p m l (insertBM kx bm r) + Tip kx' bm' + | kx' == kx -> Tip kx' (bm .|. bm') + | otherwise -> link kx (Tip kx bm) kx' t + Nil -> Tip kx bm + +-- | /O(min(n,W))/. Delete a value in the set. Returns the +-- original set when the value was not present. +delete :: Key -> IntSet -> IntSet +delete x = x `seq` deleteBM (prefixOf x) (bitmapOf x) + +-- 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 = link p1 t1 p2 t2 + where + union1 | nomatch p2 p1 m1 = link 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 = link 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 + +infixl 5 `union` + +{-------------------------------------------------------------------- + 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 + +infixl 5 `intersection` + +{-------------------------------------------------------------------- + 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 :: (Key -> 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 :: (Key -> Bool) -> IntSet -> (IntSet,IntSet) +partition predicate0 t0 = toPair $ go predicate0 t0 + where + go predicate t + = case t of + Bin p m l r + -> let (l1 :*: l2) = go predicate l + (r1 :*: r2) = go 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 :: Key -> IntSet -> (IntSet,IntSet) +split x t = + case t of + Bin _ m l r + | m < 0 -> if x >= 0 -- handle negative numbers. + then case go x l of (lt :*: gt) -> let lt' = union lt r + in lt' `seq` (lt', gt) + else case go x r of (lt :*: gt) -> let gt' = union gt l + in gt' `seq` (lt, gt') + _ -> case go x t of + (lt :*: gt) -> (lt, gt) + 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 :: Key -> 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) -> let lt' = union lt r + in lt' `seq` (lt', fnd, gt) + else case go x r of + (lt, fnd, gt) -> let gt' = union gt l + in gt' `seq` (lt, fnd, gt') + _ -> 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 = let lt = tip kx' (bm .&. lowerBitmap) + found = (bm .&. bitmapOfx') /= 0 + gt = tip kx' (bm .&. higherBitmap) + in lt `seq` found `seq` gt `seq` (lt, found, gt) + 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 (Key, 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 (Key, 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 -> (Key, 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 -> (Key, IntSet) +deleteFindMax = fromMaybe (error "deleteFindMax: empty set has no maximal element") . maxView + + +-- | /O(min(n,W))/. The minimal element of the set. +findMin :: IntSet -> Key +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 -> Key +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. Returns an empty set if the set is empty. +-- +-- Note that this is a change of behaviour for consistency with 'Data.Set.Set' – +-- versions prior to 0.5 threw an error if the 'IntSet' was already empty. +deleteMin :: IntSet -> IntSet +deleteMin = maybe Nil snd . minView + +-- | /O(min(n,W))/. Delete the maximal element. Returns an empty set if the set is empty. +-- +-- Note that this is a change of behaviour for consistency with 'Data.Set.Set' – +-- versions prior to 0.5 threw an error if the 'IntSet' was already empty. +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 :: (Key -> Key) -> 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 :: (Key -> 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 :: (Key -> 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' :: (Key -> 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 -> Key -> 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 -> Key -> 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 -> [Key] +elems + = toAscList + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} +#if __GLASGOW_HASKELL__ >= 708 +instance GHCExts.IsList IntSet where + type Item IntSet = Key + fromList = fromList + toList = toList +#endif + +-- | /O(n)/. Convert the set to a list of elements. Subject to list fusion. +toList :: IntSet -> [Key] +toList + = toAscList + +-- | /O(n)/. Convert the set to an ascending list of elements. Subject to list +-- fusion. +toAscList :: IntSet -> [Key] +toAscList = foldr (:) [] + +-- | /O(n)/. Convert the set to a descending list of elements. Subject to list +-- fusion. +toDescList :: IntSet -> [Key] +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 :: (Key -> b -> b) -> b -> IntSet -> b +foldrFB = foldr +{-# INLINE[0] foldrFB #-} +foldlFB :: (a -> Key -> 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 :: [Key] -> 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 :: [Key] -> 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 :: [Key] -> 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 (link 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 +--------------------------------------------------------------------} +{-------------------------------------------------------------------- + Link +--------------------------------------------------------------------} +link :: Prefix -> IntSet -> Prefix -> IntSet -> IntSet +link 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 link #-} + +{-------------------------------------------------------------------- + @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 +#if MIN_VERSION_base_4_7_0 +suffixBitMask = finiteBitSize (undefined::Word) - 1 +#else +suffixBitMask = bitSize (undefined::Word) - 1 +#endif +{-# 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 #-} + +{---------------------------------------------------------------------- + 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 +#if MIN_VERSION_base_4_5_0 +bitcount a x = a + popCount x +#else +bitcount a0 x0 = go a0 x0 + where go a 0 = a + go a x = go (a + 1) (x .&. (x-1)) +#endif +{-# 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 #-} + +-- | /O(1)/. Decompose a set into pieces based on the structure of the underlying +-- tree. This function is useful for consuming a set in parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that the +-- pieces returned will be in ascending order (all elements in the first submap +-- less than all elements in the second, and so on). +-- +-- Examples: +-- +-- > splitRoot (fromList [1..120]) == [fromList [1..63],fromList [64..120]] +-- > splitRoot empty == [] +-- +-- Note that the current implementation does not return more than two subsets, +-- but you should not depend on this behaviour because it can change in the +-- future without notice. Also, the current version does not continue +-- splitting all the way to individual singleton sets -- it stops at some +-- point. +splitRoot :: IntSet -> [IntSet] +splitRoot orig = + case orig of + Nil -> [] + -- NOTE: we don't currently split below Tip, but we could. + x@(Tip _ _) -> [x] + Bin _ m l r | m < 0 -> [r, l] + | otherwise -> [l, r] +{-# INLINE splitRoot #-} diff --git a/Data/Map.hs b/Data/Map.hs new file mode 100644 --- /dev/null +++ b/Data/Map.hs @@ -0,0 +1,129 @@ +{-# 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 +-- +-- /Note:/ You should use "Data.Map.Strict" instead of this module if: +-- +-- * You will eventually need all the values stored. +-- +-- * The stored values don't represent large virtual data structures +-- to be lazily computed. +-- +-- An efficient implementation of ordered maps from keys to values +-- (dictionaries). +-- +-- 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 Prelude hiding (foldr) +import Data.Map.Lazy +import qualified Data.Map.Strict as Strict + +-- | /Deprecated./ As of version 0.5, replaced by 'Data.Map.Strict.insertWith'. +-- +-- /O(log n)/. Same as 'insertWith', but the value being inserted to the map is +-- evaluated to WHNF beforehand. +-- +-- 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' = Strict.insertWith +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertWith' #-} +#else +{-# INLINE insertWith' #-} +#endif + +-- | /Deprecated./ As of version 0.5, replaced by +-- 'Data.Map.Strict.insertWithKey'. +-- +-- /O(log n)/. Same as 'insertWithKey', but the value being inserted to the map is +-- evaluated to WHNF beforehand. + +insertWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +-- We do not reuse Data.Map.Strict.insertWithKey, because it is stricter -- it +-- forces evaluation of the given value. +insertWithKey' = Strict.insertWithKey +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertWithKey' #-} +#else +{-# INLINE insertWithKey' #-} +#endif + +-- | /Deprecated./ As of version 0.5, replaced by +-- 'Data.Map.Strict.insertLookupWithKey'. +-- +-- /O(log n)/. Same as 'insertLookupWithKey', but the value being inserted to +-- the map is evaluated to WHNF beforehand. + +insertLookupWithKey' :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a + -> (Maybe a, Map k a) +-- We do not reuse Data.Map.Strict.insertLookupWithKey, because it is stricter -- it +-- forces evaluation of the given value. +insertLookupWithKey' = Strict.insertLookupWithKey +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertLookupWithKey' #-} +#else +{-# INLINE insertLookupWithKey' #-} +#endif + +-- | /Deprecated./ As of version 0.5, replaced by 'foldr'. +-- +-- /O(n)/. Fold the values in the map using the given right-associative +-- 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 = foldr +{-# INLINE fold #-} + +-- | /Deprecated./ As of version 0.4, replaced by 'foldrWithKey'. +-- +-- /O(n)/. Fold the keys and values in the map using the given right-associative +-- 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,2868 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeFamilies #-} +#endif +----------------------------------------------------------------------------- +-- | +-- 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 + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + + -- Used by the strict version + , bin + , balance + , balanced + , balanceL + , balanceR + , delta + , link + , insertMax + , merge + , glue + , trim + , trimLookupLo + , foldlStrict + , MaybeS(..) + , filterGt + , filterLt + ) where + +import Control.Applicative (Applicative(..), (<$>)) +import Control.DeepSeq (NFData(rnf)) +import Data.Bits (shiftL, shiftR) +import qualified Data.Foldable as Foldable +import Data.Monoid (Monoid(..)) +import Data.StrictPair +import Data.Traversable (Traversable(traverse)) +import Data.Typeable +import Prelude hiding (lookup, map, filter, foldr, foldl, null) + +import qualified Data.Set.Base as Set + +#if __GLASGOW_HASKELL__ +import GHC.Exts ( build ) +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts as GHCExts +#endif +import Text.Read +import Data.Data +#endif + +-- 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 + +#if __GLASGOW_HASKELL__ >= 708 +type role Map nominal representational +#endif + +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 provide limited reflection services for the sake of data abstraction. + +instance (Data k, Data a, Ord k) => Data (Map k a) where + gfoldl f z m = z fromList `f` toList m + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = mapDataType + dataCast2 f = gcast2 f + +fromListConstr :: Constr +fromListConstr = mkConstr mapDataType "fromList" [] Prefix + +mapDataType :: DataType +mapDataType = mkDataType "Data.Map.Base.Map" [fromListConstr] + +#endif + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} +-- | /O(1)/. Is the map empty? +-- +-- > Data.Map.null (empty) == True +-- > Data.Map.null (singleton 1 'a') == False + +null :: Map k a -> Bool +null Tip = True +null (Bin {}) = False +{-# INLINE null #-} + +-- | /O(1)/. The number of elements in the map. +-- +-- > size empty == 0 +-- > size (singleton 1 'a') == 1 +-- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 + +size :: Map k a -> Int +size Tip = 0 +size (Bin sz _ _ _ _) = sz +{-# INLINE size #-} + + +-- | /O(log n)/. Lookup the value at a key in the map. +-- +-- The function will return the corresponding value as @('Just' value)@, +-- or 'Nothing' if the key isn't in the map. +-- +-- An example of using @lookup@: +-- +-- > import Prelude hiding (lookup) +-- > import Data.Map +-- > +-- > employeeDept = fromList([("John","Sales"), ("Bob","IT")]) +-- > deptCountry = fromList([("IT","USA"), ("Sales","France")]) +-- > countryCurrency = fromList([("USA", "Dollar"), ("France", "Euro")]) +-- > +-- > employeeCurrency :: String -> Maybe String +-- > employeeCurrency name = do +-- > dept <- lookup name employeeDept +-- > country <- lookup dept deptCountry +-- > lookup country countryCurrency +-- > +-- > main = do +-- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John")) +-- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete")) +-- +-- The output of this program: +-- +-- > John's currency: Just "Euro" +-- > Pete's currency: Nothing +lookup :: Ord k => k -> Map k a -> Maybe a +lookup = go + where + STRICT_1_OF_2(go) + go _ Tip = Nothing + go k (Bin _ kx x l r) = case compare k kx of + LT -> go k l + GT -> go k r + EQ -> Just x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookup #-} +#else +{-# INLINE lookup #-} +#endif + +-- | /O(log n)/. Is the key a member of the map? See also 'notMember'. +-- +-- > member 5 (fromList [(5,'a'), (3,'b')]) == True +-- > member 1 (fromList [(5,'a'), (3,'b')]) == False +member :: Ord k => k -> Map k a -> Bool +member = go + where + STRICT_1_OF_2(go) + go _ Tip = False + go k (Bin _ kx _ l r) = case compare k kx of + LT -> go k l + GT -> go k r + EQ -> True +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE member #-} +#else +{-# INLINE member #-} +#endif + +infix 4 `member` + +-- | /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 + +infix 4 `notMember` + +-- | /O(log n)/. Find the value at a key. +-- Calls 'error' when the element can not be found. +find :: Ord k => k -> Map k a -> a +find = go + where + STRICT_1_OF_2(go) + go _ Tip = error "Map.!: given key is not an element in the map" + go k (Bin _ kx x l r) = case compare k kx of + LT -> go k l + GT -> go k r + EQ -> x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE find #-} +#else +{-# INLINE find #-} +#endif + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' +findWithDefault :: Ord k => a -> k -> Map k a -> a +findWithDefault = go + where + STRICT_2_OF_3(go) + go def _ Tip = def + go def k (Bin _ kx x l r) = case compare k kx of + LT -> go def k l + GT -> go def k r + EQ -> x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE findWithDefault #-} +#else +{-# INLINE findWithDefault #-} +#endif + +-- | /O(log n)/. Find largest key smaller than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing +-- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') +lookupLT :: Ord k => k -> Map k v -> Maybe (k, v) +lookupLT = goNothing + where + STRICT_1_OF_2(goNothing) + goNothing _ Tip = Nothing + goNothing k (Bin _ kx x l r) | k <= kx = goNothing k l + | otherwise = goJust k kx x r + + STRICT_1_OF_4(goJust) + goJust _ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) | k <= kx = goJust k kx' x' l + | otherwise = goJust k kx x r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupLT #-} +#else +{-# INLINE lookupLT #-} +#endif + +-- | /O(log n)/. Find smallest key greater than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +-- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing +lookupGT :: Ord k => k -> Map k v -> Maybe (k, v) +lookupGT = goNothing + where + STRICT_1_OF_2(goNothing) + goNothing _ Tip = Nothing + goNothing k (Bin _ kx x l r) | k < kx = goJust k kx x l + | otherwise = goNothing k r + + STRICT_1_OF_4(goJust) + goJust _ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) | k < kx = goJust k kx x l + | otherwise = goJust k kx' x' r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupGT #-} +#else +{-# INLINE lookupGT #-} +#endif + +-- | /O(log n)/. Find largest key smaller or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupLE 2 (fromList [(3,'a'), (5,'b')]) == Nothing +-- > lookupLE 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') +-- > lookupLE 5 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +lookupLE :: Ord k => k -> Map k v -> Maybe (k, v) +lookupLE = goNothing + where + STRICT_1_OF_2(goNothing) + goNothing _ Tip = Nothing + goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goNothing k l + EQ -> Just (kx, x) + GT -> goJust k kx x r + + STRICT_1_OF_4(goJust) + goJust _ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l + EQ -> Just (kx, x) + GT -> goJust k kx x r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupLE #-} +#else +{-# INLINE lookupLE #-} +#endif + +-- | /O(log n)/. Find smallest key greater or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupGE 3 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') +-- > lookupGE 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +-- > lookupGE 6 (fromList [(3,'a'), (5,'b')]) == Nothing +lookupGE :: Ord k => k -> Map k v -> Maybe (k, v) +lookupGE = goNothing + where + STRICT_1_OF_2(goNothing) + goNothing _ Tip = Nothing + goNothing k (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l + EQ -> Just (kx, x) + GT -> goNothing k r + + STRICT_1_OF_4(goJust) + goJust _ kx' x' Tip = Just (kx', x') + goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l + EQ -> Just (kx, x) + GT -> goJust k kx' x' r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupGE #-} +#else +{-# INLINE lookupGE #-} +#endif + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} +-- | /O(1)/. The empty map. +-- +-- > empty == fromList [] +-- > size empty == 0 + +empty :: Map k a +empty = Tip +{-# INLINE empty #-} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 + +singleton :: k -> a -> Map k a +singleton k x = Bin 1 k x Tip Tip +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insertion +--------------------------------------------------------------------} +-- | /O(log n)/. Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value. 'insert' is equivalent to +-- @'insertWith' 'const'@. +-- +-- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')] +-- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')] +-- > insert 5 'x' empty == singleton 5 'x' + +-- See Note: Type of local 'go' function +insert :: Ord k => k -> a -> Map k a -> Map k a +insert = go + where + go :: Ord k => k -> a -> Map k a -> Map k a + STRICT_1_OF_3(go) + go kx x Tip = singleton kx x + go kx x (Bin sz ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go kx x l) r + GT -> balanceR ky y l (go kx x r) + EQ -> Bin sz kx x l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insert #-} +#else +{-# INLINE insert #-} +#endif + +-- Insert a new key and value in the map if it is not already present. +-- Used by `union`. + +-- See Note: Type of local 'go' function +insertR :: Ord k => k -> a -> Map k a -> Map k a +insertR = go + where + go :: Ord k => k -> a -> Map k a -> Map k a + STRICT_1_OF_3(go) + go kx x Tip = singleton kx x + go kx x t@(Bin _ ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go kx x l) r + GT -> balanceR ky y l (go kx x r) + EQ -> t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertR #-} +#else +{-# INLINE insertR #-} +#endif + +-- | /O(log n)/. Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key, f new_value old_value)@. +-- +-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] +-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" + +insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWith f = insertWithKey (\_ x' y' -> f x' y') +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertWith #-} +#else +{-# INLINE insertWith #-} +#endif + +-- | /O(log n)/. Insert with a function, combining key, new value and old value. +-- @'insertWithKey' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key,f key new_value old_value)@. +-- Note that the key passed to f is the same key passed to 'insertWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] +-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" + +-- See Note: Type of local 'go' function +insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKey = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a + STRICT_2_OF_4(go) + go _ kx x Tip = singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> Bin sy kx (f kx x y) l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertWithKey #-} +#else +{-# INLINE insertWithKey #-} +#endif + +-- | /O(log n)/. Combines insert operation with old value retrieval. +-- The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- and the second element equal to (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) +-- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) +-- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) + +-- See Note: Type of local 'go' function +insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a + -> (Maybe a, Map k a) +insertLookupWithKey = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> (Maybe a, Map k a) + STRICT_2_OF_4(go) + go _ kx x Tip = (Nothing, singleton kx x) + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> let (found, l') = go f kx x l + in (found, balanceL ky y l' r) + GT -> let (found, r') = go f kx x r + in (found, balanceR ky y l r') + EQ -> (Just y, Bin sy kx (f kx x y) l r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertLookupWithKey #-} +#else +{-# INLINE insertLookupWithKey #-} +#endif + +{-------------------------------------------------------------------- + Deletion +--------------------------------------------------------------------} +-- | /O(log n)/. Delete a key and its value from the map. When the key is not +-- a member of the map, the original map is returned. +-- +-- > delete 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > delete 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > delete 5 empty == empty + +-- See Note: Type of local 'go' function +delete :: Ord k => k -> Map k a -> Map k a +delete = go + where + go :: Ord k => k -> Map k a -> Map k a + STRICT_1_OF_2(go) + go _ Tip = Tip + go k (Bin _ kx x l r) = + case compare k kx of + LT -> balanceR kx x (go k l) r + GT -> balanceL kx x l (go k r) + EQ -> glue l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE delete #-} +#else +{-# INLINE delete #-} +#endif + +-- | /O(log n)/. Update a value at a specific key with the result of the provided function. +-- When the key is not +-- a member of the map, the original map is returned. +-- +-- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > adjust ("new " ++) 7 empty == empty + +adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a +adjust f = adjustWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE adjust #-} +#else +{-# INLINE adjust #-} +#endif + +-- | /O(log n)/. Adjust a value at a specific key. When the key is not +-- a member of the map, the original map is returned. +-- +-- > let f key x = (show key) ++ ":new " ++ x +-- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > adjustWithKey f 7 empty == empty + +adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a +adjustWithKey f = updateWithKey (\k' x' -> Just (f k' x')) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE adjustWithKey #-} +#else +{-# INLINE adjustWithKey #-} +#endif + +-- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a +update f = updateWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE update #-} +#else +{-# INLINE update #-} +#endif + +-- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the +-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', +-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound +-- to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +-- See Note: Type of local 'go' function +updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a +updateWithKey = go + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a + STRICT_2_OF_3(go) + go _ _ Tip = Tip + go f k(Bin sx kx x l r) = + case compare k kx of + LT -> balanceR kx x (go f k l) r + GT -> balanceL kx x l (go f k r) + EQ -> case f kx x of + Just x' -> Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE updateWithKey #-} +#else +{-# INLINE updateWithKey #-} +#endif + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") + +-- See Note: Type of local 'go' function +updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) +updateLookupWithKey = go + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) + STRICT_2_OF_3(go) + go _ _ Tip = (Nothing,Tip) + go f k (Bin sx kx x l r) = + case compare k kx of + LT -> let (found,l') = go f k l in (found,balanceR kx x l' r) + GT -> let (found,r') = go f k r in (found,balanceL kx x l r') + EQ -> case f kx x of + Just x' -> (Just x',Bin sx kx x' l r) + Nothing -> (Just x,glue l r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE updateLookupWithKey #-} +#else +{-# INLINE updateLookupWithKey #-} +#endif + +-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. +-- 'alter' can be used to insert, delete, or update a value in a 'Map'. +-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] + +-- See Note: Type of local 'go' function +alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +alter = go + where + go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a + STRICT_2_OF_3(go) + go f k Tip = case f Nothing of + Nothing -> Tip + Just x -> singleton k x + + go f k (Bin sx kx x l r) = case compare k kx of + LT -> balance kx x (go f k l) r + GT -> balance kx x l (go f k r) + EQ -> case f (Just x) of + Just x' -> Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE alter #-} +#else +{-# INLINE alter #-} +#endif + +{-------------------------------------------------------------------- + Indexing +--------------------------------------------------------------------} +-- | /O(log n)/. Return the /index/ of a key, which is its zero-based index in +-- the sequence sorted by keys. The index is a number from /0/ up to, but not +-- including, the 'size' of the map. Calls 'error' when the key is not +-- a 'member' of the map. +-- +-- > findIndex 2 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map +-- > findIndex 3 (fromList [(5,"a"), (3,"b")]) == 0 +-- > findIndex 5 (fromList [(5,"a"), (3,"b")]) == 1 +-- > findIndex 6 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map + +-- See Note: Type of local 'go' function +findIndex :: Ord k => k -> Map k a -> Int +findIndex = go 0 + where + go :: Ord k => Int -> k -> Map k a -> Int + STRICT_1_OF_3(go) + STRICT_2_OF_3(go) + go _ _ Tip = error "Map.findIndex: element is not in the map" + go idx k (Bin _ kx _ l r) = case compare k kx of + LT -> go idx k l + GT -> go (idx + size l + 1) k r + EQ -> idx + size l +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE findIndex #-} +#endif + +-- | /O(log n)/. Lookup the /index/ of a key, which is its zero-based index in +-- the sequence sorted by keys. The index is a number from /0/ up to, but not +-- including, the 'size' of the map. +-- +-- > isJust (lookupIndex 2 (fromList [(5,"a"), (3,"b")])) == False +-- > fromJust (lookupIndex 3 (fromList [(5,"a"), (3,"b")])) == 0 +-- > fromJust (lookupIndex 5 (fromList [(5,"a"), (3,"b")])) == 1 +-- > isJust (lookupIndex 6 (fromList [(5,"a"), (3,"b")])) == False + +-- See Note: Type of local 'go' function +lookupIndex :: Ord k => k -> Map k a -> Maybe Int +lookupIndex = go 0 + where + go :: Ord k => Int -> k -> Map k a -> Maybe Int + STRICT_1_OF_3(go) + STRICT_2_OF_3(go) + go _ _ Tip = Nothing + go idx k (Bin _ kx _ l r) = case compare k kx of + LT -> go idx k l + GT -> go (idx + size l + 1) k r + EQ -> Just $! idx + size l +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupIndex #-} +#endif + +-- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based +-- index in the sequence sorted by keys. If the /index/ is out of range (less +-- than zero, greater or equal to 'size' of the map), 'error' is called. +-- +-- > elemAt 0 (fromList [(5,"a"), (3,"b")]) == (3,"b") +-- > elemAt 1 (fromList [(5,"a"), (3,"b")]) == (5, "a") +-- > elemAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range + +elemAt :: Int -> Map k a -> (k,a) +STRICT_1_OF_2(elemAt) +elemAt _ Tip = error "Map.elemAt: index out of range" +elemAt i (Bin _ kx x l r) + = case compare i sizeL of + LT -> elemAt i l + GT -> elemAt (i-sizeL-1) r + EQ -> (kx,x) + where + sizeL = size l + +-- | /O(log n)/. Update the element at /index/, i.e. by its zero-based index in +-- the sequence sorted by keys. If the /index/ is out of range (less than zero, +-- greater or equal to 'size' of the map), 'error' is called. +-- +-- > updateAt (\ _ _ -> Just "x") 0 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "x"), (5, "a")] +-- > updateAt (\ _ _ -> Just "x") 1 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "x")] +-- > updateAt (\ _ _ -> Just "x") 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > updateAt (\ _ _ -> Just "x") (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > updateAt (\_ _ -> Nothing) 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +-- > updateAt (\_ _ -> Nothing) 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > updateAt (\_ _ -> Nothing) 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > updateAt (\_ _ -> Nothing) (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range + +updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a +updateAt f i t = i `seq` + case t of + Tip -> error "Map.updateAt: index out of range" + Bin sx kx x l r -> case compare i sizeL of + LT -> balanceR kx x (updateAt f i l) r + GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) + EQ -> case f kx x of + Just x' -> Bin sx kx x' l r + Nothing -> glue l r + where + sizeL = size l + +-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in +-- the sequence sorted by keys. If the /index/ is out of range (less than zero, +-- greater or equal to 'size' of the map), 'error' is called. +-- +-- > deleteAt 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +-- > deleteAt 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > deleteAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > deleteAt (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range + +deleteAt :: Int -> Map k a -> Map k a +deleteAt i t = i `seq` + case t of + Tip -> error "Map.deleteAt: index out of range" + Bin _ kx x l r -> case compare i sizeL of + LT -> balanceR kx x (deleteAt i l) r + GT -> balanceL kx x l (deleteAt (i-sizeL-1) r) + EQ -> glue l r + where + sizeL = size l + + +{-------------------------------------------------------------------- + Minimal, Maximal +--------------------------------------------------------------------} +-- | /O(log n)/. The minimal key of the map. Calls 'error' if the map is empty. +-- +-- > findMin (fromList [(5,"a"), (3,"b")]) == (3,"b") +-- > findMin empty Error: empty map has no minimal element + +findMin :: Map k a -> (k,a) +findMin (Bin _ kx x Tip _) = (kx,x) +findMin (Bin _ _ _ l _) = findMin l +findMin Tip = error "Map.findMin: empty map has no minimal element" + +-- | /O(log n)/. The maximal key of the map. Calls 'error' if the map is empty. +-- +-- > findMax (fromList [(5,"a"), (3,"b")]) == (5,"a") +-- > findMax empty Error: empty map has no maximal element + +findMax :: Map k a -> (k,a) +findMax (Bin _ kx x _ Tip) = (kx,x) +findMax (Bin _ _ _ _ r) = findMax r +findMax Tip = error "Map.findMax: empty map has no maximal element" + +-- | /O(log n)/. Delete the minimal key. Returns an empty map if the map is empty. +-- +-- > deleteMin (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(5,"a"), (7,"c")] +-- > deleteMin empty == empty + +deleteMin :: Map k a -> Map k a +deleteMin (Bin _ _ _ Tip r) = r +deleteMin (Bin _ kx x l r) = balanceR kx x (deleteMin l) r +deleteMin Tip = Tip + +-- | /O(log n)/. Delete the maximal key. Returns an empty map if the map is empty. +-- +-- > deleteMax (fromList [(5,"a"), (3,"b"), (7,"c")]) == fromList [(3,"b"), (5,"a")] +-- > deleteMax empty == empty + +deleteMax :: Map k a -> Map k a +deleteMax (Bin _ _ _ l Tip) = l +deleteMax (Bin _ kx x l r) = balanceL kx x l (deleteMax r) +deleteMax Tip = Tip + +-- | /O(log n)/. Update the value at the minimal key. +-- +-- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")] +-- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +updateMin :: (a -> Maybe a) -> Map k a -> Map k a +updateMin f m + = updateMinWithKey (\_ x -> f x) m + +-- | /O(log n)/. Update the value at the maximal key. +-- +-- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")] +-- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" + +updateMax :: (a -> Maybe a) -> Map k a -> Map k a +updateMax f m + = updateMaxWithKey (\_ x -> f x) m + + +-- | /O(log n)/. Update the value at the minimal key. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMinWithKey _ Tip = Tip +updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of + Nothing -> r + Just x' -> Bin sx kx x' Tip r +updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r + +-- | /O(log n)/. Update the value at the maximal key. +-- +-- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")] +-- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" + +updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMaxWithKey _ Tip = Tip +updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of + Nothing -> l + Just x' -> Bin sx kx x' l Tip +updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) + +-- | /O(log n)/. Retrieves the minimal (key,value) pair of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +-- +-- > minViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((3,"b"), singleton 5 "a") +-- > minViewWithKey empty == Nothing + +minViewWithKey :: Map k a -> Maybe ((k,a), Map k a) +minViewWithKey Tip = Nothing +minViewWithKey x = Just (deleteFindMin x) + +-- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and +-- the map stripped of that element, or 'Nothing' if passed an empty map. +-- +-- > maxViewWithKey (fromList [(5,"a"), (3,"b")]) == Just ((5,"a"), singleton 3 "b") +-- > maxViewWithKey empty == Nothing + +maxViewWithKey :: Map k a -> Maybe ((k,a), Map k a) +maxViewWithKey Tip = Nothing +maxViewWithKey x = Just (deleteFindMax x) + +-- | /O(log n)/. Retrieves the value associated with minimal key of the +-- map, and the map stripped of that element, or 'Nothing' if passed an +-- empty map. +-- +-- > minView (fromList [(5,"a"), (3,"b")]) == Just ("b", singleton 5 "a") +-- > minView empty == Nothing + +minView :: Map k a -> Maybe (a, Map k a) +minView Tip = Nothing +minView x = Just (first snd $ deleteFindMin x) + +-- | /O(log n)/. Retrieves the value associated with maximal key of the +-- map, and the map stripped of that element, or 'Nothing' if passed an +-- +-- > maxView (fromList [(5,"a"), (3,"b")]) == Just ("a", singleton 3 "b") +-- > maxView empty == Nothing + +maxView :: Map k a -> Maybe (a, Map k a) +maxView Tip = Nothing +maxView x = Just (first snd $ deleteFindMax x) + +-- Update the 1st component of a tuple (special case of Control.Arrow.first) +first :: (a -> b) -> (a,c) -> (b,c) +first f (x,y) = (f x, y) + +{-------------------------------------------------------------------- + Union. +--------------------------------------------------------------------} +-- | The union of a list of maps: +-- (@'unions' == 'Prelude.foldl' 'union' 'empty'@). +-- +-- > unions [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "b"), (5, "a"), (7, "C")] +-- > unions [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] +-- > == fromList [(3, "B3"), (5, "A3"), (7, "C")] + +unions :: Ord k => [Map k a] -> Map k a +unions ts + = foldlStrict union empty ts +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unions #-} +#endif + +-- | The union of a list of maps, with a combining operation: +-- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a +unionsWith f ts + = foldlStrict (unionWith f) empty ts +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unionsWith #-} +#endif + +-- | /O(n+m)/. +-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@. +-- It prefers @t1@ when duplicate keys are encountered, +-- i.e. (@'union' == 'unionWith' 'const'@). +-- The implementation uses the efficient /hedge-union/ algorithm. +-- +-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] + +union :: Ord k => Map k a -> Map k a -> Map k a +union Tip t2 = t2 +union t1 Tip = t1 +union t1 t2 = hedgeUnion NothingS NothingS t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE union #-} +#endif + +infixl 5 `union` + +-- left-biased hedge union +hedgeUnion :: Ord a => MaybeS a -> MaybeS a -> Map a b -> Map a b -> Map a b +hedgeUnion _ _ t1 Tip = t1 +hedgeUnion blo bhi Tip (Bin _ kx x l r) = link kx x (filterGt blo l) (filterLt bhi r) +hedgeUnion _ _ t1 (Bin _ kx x Tip Tip) = insertR kx x t1 -- According to benchmarks, this special case increases + -- performance up to 30%. It does not help in difference or intersection. +hedgeUnion blo bhi (Bin _ kx x l r) t2 = link kx x (hedgeUnion blo bmi l (trim blo bmi t2)) + (hedgeUnion bmi bhi r (trim bmi bhi t2)) + where bmi = JustS kx +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE hedgeUnion #-} +#endif + +{-------------------------------------------------------------------- + Union with a combining function +--------------------------------------------------------------------} +-- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] + +unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWith f m1 m2 + = unionWithKey (\_ x y -> f x y) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unionWith #-} +#endif + +-- | /O(n+m)/. +-- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] + +unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) id id t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unionWithKey #-} +#endif + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} +-- | /O(n+m)/. Difference of two maps. +-- Return elements of the first map not existing in the second map. +-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" + +difference :: Ord k => Map k a -> Map k b -> Map k a +difference Tip _ = Tip +difference t1 Tip = t1 +difference t1 t2 = hedgeDiff NothingS NothingS t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE difference #-} +#endif + +hedgeDiff :: Ord a => MaybeS a -> MaybeS a -> Map a b -> Map a c -> Map a b +hedgeDiff _ _ Tip _ = Tip +hedgeDiff blo bhi (Bin _ kx x l r) Tip = link kx x (filterGt blo l) (filterLt bhi r) +hedgeDiff blo bhi t (Bin _ kx _ l r) = merge (hedgeDiff blo bmi (trim blo bmi t) l) + (hedgeDiff bmi bhi (trim bmi bhi t) r) + where bmi = JustS kx +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE hedgeDiff #-} +#endif + +-- | /O(n+m)/. Difference with a combining function. +-- When two equal keys are +-- encountered, the combining function is applied to the values of these keys. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) +-- > == singleton 3 "b:B" + +differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWith f m1 m2 + = differenceWithKey (\_ x y -> f x y) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE differenceWith #-} +#endif + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) +-- > == singleton 3 "3:b|B" + +differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWithKey f t1 t2 = mergeWithKey f id (const Tip) t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE differenceWithKey #-} +#endif + + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} +-- | /O(n+m)/. Intersection of two maps. +-- Return data in the first map for the keys existing in both maps. +-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). +-- The implementation uses an efficient /hedge/ algorithm comparable with +-- /hedge-union/. +-- +-- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" + +intersection :: Ord k => Map k a -> Map k b -> Map k a +intersection Tip _ = Tip +intersection _ Tip = Tip +intersection t1 t2 = hedgeInt NothingS NothingS t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE intersection #-} +#endif + +infixl 5 `intersection` + +hedgeInt :: Ord k => MaybeS k -> MaybeS k -> Map k a -> Map k b -> Map k a +hedgeInt _ _ _ Tip = Tip +hedgeInt _ _ Tip _ = Tip +hedgeInt blo bhi (Bin _ kx x l r) t2 = let l' = hedgeInt blo bmi l (trim blo bmi t2) + r' = hedgeInt bmi bhi r (trim bmi bhi t2) + in if kx `member` t2 then link kx x l' r' else merge l' r' + where bmi = JustS kx +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE hedgeInt #-} +#endif + +-- | /O(n+m)/. Intersection with a combining function. The implementation uses +-- an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" + +intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWith f m1 m2 + = intersectionWithKey (\_ x y -> f x y) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE intersectionWith #-} +#endif + +-- | /O(n+m)/. Intersection with a combining function. The implementation uses +-- an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" + + +intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) (const Tip) (const Tip) t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE intersectionWithKey #-} +#endif + + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | /O(n+m)/. A high-performance universal combining function. This function +-- is used to define 'unionWith', 'unionWithKey', 'differenceWith', +-- 'differenceWithKey', 'intersectionWith', 'intersectionWithKey' and can be +-- used to define other custom combine functions. +-- +-- Please make sure you know what is going on when using 'mergeWithKey', +-- otherwise you can be surprised by unexpected code growth or even +-- corruption of the data structure. +-- +-- When 'mergeWithKey' is given three arguments, it is inlined to the call +-- site. You should therefore use 'mergeWithKey' only to define your custom +-- combining functions. For example, you could define 'unionWithKey', +-- 'differenceWithKey' and 'intersectionWithKey' as +-- +-- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 +-- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 +-- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 +-- +-- When calling @'mergeWithKey' combine only1 only2@, a function combining two +-- 'IntMap's is created, such that +-- +-- * if a key is present in both maps, it is passed with both corresponding +-- values to the @combine@ function. Depending on the result, the key is either +-- present in the result with specified value, or is left out; +-- +-- * a nonempty subtree present only in the first map is passed to @only1@ and +-- the output is added to the result; +-- +-- * a nonempty subtree present only in the second map is passed to @only2@ and +-- the output is added to the result. +-- +-- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. +-- The values can be modified arbitrarily. Most common variants of @only1@ and +-- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or +-- @'filterWithKey' f@ could be used for any @f@. + +mergeWithKey :: Ord k => (k -> a -> b -> Maybe c) -> (Map k a -> Map k c) -> (Map k b -> Map k c) + -> Map k a -> Map k b -> Map k c +mergeWithKey f g1 g2 = go + where + go Tip t2 = g2 t2 + go t1 Tip = g1 t1 + go t1 t2 = hedgeMerge NothingS NothingS t1 t2 + + hedgeMerge _ _ t1 Tip = g1 t1 + hedgeMerge blo bhi Tip (Bin _ kx x l r) = g2 $ link kx x (filterGt blo l) (filterLt bhi r) + hedgeMerge blo bhi (Bin _ kx x l r) t2 = let l' = hedgeMerge blo bmi l (trim blo bmi t2) + (found, trim_t2) = trimLookupLo kx bhi t2 + r' = hedgeMerge bmi bhi r trim_t2 + in case found of + Nothing -> case g1 (singleton kx x) of + Tip -> merge l' r' + (Bin _ _ x' Tip Tip) -> link kx x' l' r' + _ -> error "mergeWithKey: Given function only1 does not fulfil required conditions (see documentation)" + Just x2 -> case f kx x x2 of + Nothing -> merge l' r' + Just x' -> link kx x' l' r' + where bmi = JustS kx +{-# INLINE mergeWithKey #-} + +{-------------------------------------------------------------------- + Submap +--------------------------------------------------------------------} +-- | /O(n+m)/. +-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +-- +isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool +isSubmapOf m1 m2 = isSubmapOfBy (==) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isSubmapOf #-} +#endif + +{- | /O(n+m)/. + The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if + all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when + applied to their respective values. For example, the following + expressions are all 'True': + + > isSubmapOfBy (==) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (<=) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1),('b',2)]) + + But the following are all 'False': + + > isSubmapOfBy (==) (fromList [('a',2)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (<) (fromList [('a',1)]) (fromList [('a',1),('b',2)]) + > isSubmapOfBy (==) (fromList [('a',1),('b',2)]) (fromList [('a',1)]) + + +-} +isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool +isSubmapOfBy f t1 t2 + = (size t1 <= size t2) && (submap' f t1 t2) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isSubmapOfBy #-} +#endif + +submap' :: Ord a => (b -> c -> Bool) -> Map a b -> Map a c -> Bool +submap' _ Tip _ = True +submap' _ _ Tip = False +submap' f (Bin _ kx x l r) t + = case found of + Nothing -> False + Just y -> f x y && submap' f l lt && submap' f r gt + where + (lt,found,gt) = splitLookup kx t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE submap' #-} +#endif + +-- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). +-- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). +isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool +isProperSubmapOf m1 m2 + = isProperSubmapOfBy (==) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isProperSubmapOf #-} +#endif + +{- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). + The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when + @m1@ and @m2@ are not equal, + all keys in @m1@ are in @m2@, and when @f@ returns 'True' when + applied to their respective values. For example, the following + expressions are all 'True': + + > isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + + But the following are all 'False': + + > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) + > isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) + > isProperSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + + +-} +isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool +isProperSubmapOfBy f t1 t2 + = (size t1 < size t2) && (submap' f t1 t2) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE isProperSubmapOfBy #-} +#endif + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} +-- | /O(n)/. Filter all values that satisfy the predicate. +-- +-- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty +-- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty + +filter :: (a -> Bool) -> Map k a -> Map k a +filter p m + = filterWithKey (\_ x -> p x) m + +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +filterWithKey :: (k -> a -> Bool) -> Map k a -> Map k a +filterWithKey _ Tip = Tip +filterWithKey p (Bin _ kx x l r) + | p kx x = link kx x (filterWithKey p l) (filterWithKey p r) + | otherwise = merge (filterWithKey p l) (filterWithKey p r) + +-- | /O(n)/. Partition the map according to a predicate. The first +-- map contains all elements that satisfy the predicate, the second all +-- elements that fail the predicate. See also 'split'. +-- +-- > partition (> "a") (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") +-- > partition (< "x") (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partition (> "x") (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) + +partition :: (a -> Bool) -> Map k a -> (Map k a,Map k a) +partition p m + = partitionWithKey (\_ x -> p x) m + +-- | /O(n)/. Partition the map according to a predicate. The first +-- map contains all elements that satisfy the predicate, the second all +-- elements that fail the predicate. See also 'split'. +-- +-- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b") +-- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) +-- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) + +partitionWithKey :: (k -> a -> Bool) -> Map k a -> (Map k a,Map k a) +partitionWithKey p0 t0 = toPair $ go p0 t0 + where + go _ Tip = (Tip :*: Tip) + go p (Bin _ kx x l r) + | p kx x = link kx x l1 r1 :*: merge l2 r2 + | otherwise = merge l1 r1 :*: link kx x l2 r2 + where + (l1 :*: l2) = go p l + (r1 :*: r2) = go p r + +-- | /O(n)/. Map values and collect the 'Just' results. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a" + +mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b +mapMaybe f = mapMaybeWithKey (\_ x -> f x) + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +-- +-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing +-- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" + +mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b +mapMaybeWithKey _ Tip = Tip +mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of + Just y -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) + Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r) + +-- | /O(n)/. Map values and separate the 'Left' and 'Right' results. +-- +-- > let f a = if a < "c" then Left a else Right a +-- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")]) +-- > +-- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) + +mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEither f m + = mapEitherWithKey (\_ x -> f x) m + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- +-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) +-- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")]) +-- > +-- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) + +mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEitherWithKey f0 t0 = toPair $ go f0 t0 + where + go _ Tip = (Tip :*: Tip) + go f (Bin _ kx x l r) = case f kx x of + Left y -> link kx y l1 r1 :*: merge l2 r2 + Right z -> merge l1 r1 :*: link kx z l2 r2 + where + (l1 :*: l2) = go f l + (r1 :*: r2) = go f r + +{-------------------------------------------------------------------- + Mapping +--------------------------------------------------------------------} +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] + +map :: (a -> b) -> Map k a -> Map k b +map _ Tip = Tip +map f (Bin sx kx x l r) = Bin sx kx (f x) (map f l) (map f r) + +-- | /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 +traverseWithKey :: Applicative t => (k -> a -> t b) -> Map k a -> t (Map k b) +traverseWithKey f = go + where + go Tip = pure Tip + go (Bin 1 k v _ _) = (\v' -> Bin 1 k v' Tip Tip) <$> f k v + go (Bin s k v l r) = flip (Bin s k) <$> go l <*> f k v <*> go r +{-# INLINE traverseWithKey #-} + +-- | /O(n)/. The function 'mapAccum' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a b = (a ++ b, b ++ "X") +-- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) + +mapAccum :: (a -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccum f a m + = mapAccumWithKey (\a' _ x' -> f a' x') a m + +-- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") +-- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) + +mapAccumWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumWithKey f a t + = mapAccumL f a t + +-- | /O(n)/. The function 'mapAccumL' threads an accumulating +-- argument through the map in ascending order of keys. +mapAccumL :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumL _ a Tip = (a,Tip) +mapAccumL f a (Bin sx kx x l r) = + let (a1,l') = mapAccumL f a l + (a2,x') = f a1 kx x + (a3,r') = mapAccumL f a2 r + in (a3,Bin sx kx x' l' r') + +-- | /O(n)/. The function 'mapAccumR' threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey :: (a -> k -> b -> (a,c)) -> a -> Map k b -> (a,Map k c) +mapAccumRWithKey _ a Tip = (a,Tip) +mapAccumRWithKey f a (Bin sx kx x l r) = + let (a1,r') = mapAccumRWithKey f a r + (a2,x') = f a1 kx x + (a3,l') = mapAccumRWithKey f a2 l + in (a3,Bin sx kx x' l' r') + +-- | /O(n*log n)/. +-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the value at the greatest of the +-- original keys is retained. +-- +-- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")] +-- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c" +-- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c" + +mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a +mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) [] +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE mapKeys #-} +#endif + +-- | /O(n*log n)/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" + +mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a +mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE mapKeysWith #-} +#endif + + +-- | /O(n)/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has better performance than 'mapKeys'. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] +-- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")])) == True +-- > valid (mapKeysMonotonic (\ _ -> 1) (fromList [(5,"a"), (3,"b")])) == False + +mapKeysMonotonic :: (k1->k2) -> Map k1 a -> Map k2 a +mapKeysMonotonic _ Tip = Tip +mapKeysMonotonic f (Bin sz k x l r) = + Bin sz (f k) x (mapKeysMonotonic f l) (mapKeysMonotonic f r) + +{-------------------------------------------------------------------- + Folds +--------------------------------------------------------------------} + +-- | /O(n)/. Fold the values in the map using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@. +-- +-- For example, +-- +-- > elems map = foldr (:) [] map +-- +-- > let f a len = len + (length a) +-- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +foldr :: (a -> b -> b) -> b -> Map k a -> b +foldr f z = go z + where + go z' Tip = z' + go z' (Bin _ _ x l r) = go (f x (go z' r)) l +{-# INLINE foldr #-} + +-- | /O(n)/. A strict version of 'foldr'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr' :: (a -> b -> b) -> b -> Map k a -> b +foldr' f z = go z + where + STRICT_1_OF_2(go) + go z' Tip = z' + go z' (Bin _ _ x l r) = go (f x (go z' r)) l +{-# INLINE foldr' #-} + +-- | /O(n)/. Fold the values in the map using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@. +-- +-- For example, +-- +-- > elems = reverse . foldl (flip (:)) [] +-- +-- > let f len a = len + (length a) +-- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +foldl :: (a -> b -> a) -> a -> Map k b -> a +foldl f z = go z + where + go z' Tip = z' + go z' (Bin _ _ x l r) = go (f (go z' l) x) r +{-# INLINE foldl #-} + +-- | /O(n)/. A strict version of 'foldl'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl' :: (a -> b -> a) -> a -> Map k b -> a +foldl' f z = go z + where + STRICT_1_OF_2(go) + go z' Tip = z' + go z' (Bin _ _ x l r) = go (f (go z' l) x) r +{-# INLINE foldl' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given right-associative +-- binary operator, such that +-- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. +-- +-- For example, +-- +-- > keys map = foldrWithKey (\k x ks -> k:ks) [] map +-- +-- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" +-- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)" +foldrWithKey :: (k -> a -> b -> b) -> b -> Map k a -> b +foldrWithKey f z = go z + where + go z' Tip = z' + go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l +{-# INLINE foldrWithKey #-} + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldrWithKey' :: (k -> a -> b -> b) -> b -> Map k a -> b +foldrWithKey' f z = go z + where + STRICT_1_OF_2(go) + go z' Tip = z' + go z' (Bin _ kx x l r) = go (f kx x (go z' r)) l +{-# INLINE foldrWithKey' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given left-associative +-- binary operator, such that +-- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@. +-- +-- For example, +-- +-- > keys = reverse . foldlWithKey (\ks k x -> k:ks) [] +-- +-- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" +-- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)" +foldlWithKey :: (a -> k -> b -> a) -> a -> Map k b -> a +foldlWithKey f z = go z + where + go z' Tip = z' + go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r +{-# INLINE foldlWithKey #-} + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldlWithKey' :: (a -> k -> b -> a) -> a -> Map k b -> a +foldlWithKey' f z = go z + where + STRICT_1_OF_2(go) + go z' Tip = z' + go z' (Bin _ kx x l r) = go (f (go z' l) kx x) r +{-# INLINE foldlWithKey' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given monoid, such that +-- +-- @'foldMapWithKey' f = 'Prelude.fold' . 'mapWithKey' f@ +-- +-- This can be an asymptotically faster than 'foldrWithKey' or 'foldlWithKey' for some monoids. +foldMapWithKey :: Monoid m => (k -> a -> m) -> Map k a -> m +foldMapWithKey f = go + where + go Tip = mempty + go (Bin 1 k v _ _) = f k v + go (Bin _ k v l r) = go l `mappend` (f k v `mappend` go r) +{-# INLINE foldMapWithKey #-} + +{-------------------------------------------------------------------- + List variations +--------------------------------------------------------------------} +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- Subject to list fusion. +-- +-- > elems (fromList [(5,"a"), (3,"b")]) == ["b","a"] +-- > elems empty == [] + +elems :: Map k a -> [a] +elems = foldr (:) [] + +-- | /O(n)/. Return all keys of the map in ascending order. Subject to list +-- fusion. +-- +-- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] +-- > keys empty == [] + +keys :: Map k a -> [k] +keys = foldrWithKey (\k _ ks -> k : ks) [] + +-- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map +-- in ascending key order. Subject to list fusion. +-- +-- > assocs (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > assocs empty == [] + +assocs :: Map k a -> [(k,a)] +assocs m + = toAscList m + +-- | /O(n)/. The set of all keys of the map. +-- +-- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.Set.fromList [3,5] +-- > keysSet empty == Data.Set.empty + +keysSet :: Map k a -> Set.Set k +keysSet Tip = Set.Tip +keysSet (Bin sz kx _ l r) = Set.Bin sz kx (keysSet l) (keysSet r) + +-- | /O(n)/. Build a map from a set of keys and a function which for each key +-- computes its value. +-- +-- > fromSet (\k -> replicate k 'a') (Data.Set.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] +-- > fromSet undefined Data.Set.empty == empty + +fromSet :: (k -> a) -> Set.Set k -> Map k a +fromSet _ Set.Tip = Tip +fromSet f (Set.Bin sz x l r) = Bin sz x (f x) (fromSet f l) (fromSet f r) + +{-------------------------------------------------------------------- + Lists + use [foldlStrict] to reduce demand on the control-stack +--------------------------------------------------------------------} +#if __GLASGOW_HASKELL__ >= 708 +instance (Ord k) => GHCExts.IsList (Map k v) where + type Item (Map k v) = (k,v) + fromList = fromList + toList = toList +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs. See also 'fromAscList'. +-- If the list contains more than one value for the same key, the last value +-- for the key is retained. +-- +-- If the keys of the list are ordered, linear-time implementation is used, +-- with the performance equal to 'fromDistinctAscList'. +-- +-- > fromList [] == empty +-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")] +-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")] + +-- For some reason, when 'singleton' is used in fromList or in +-- create, it is not inlined, so we inline it manually. +fromList :: Ord k => [(k,a)] -> Map k a +fromList [] = Tip +fromList [(kx, x)] = Bin 1 kx x Tip Tip +fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = fromList' (Bin 1 kx0 x0 Tip Tip) xs0 + | otherwise = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + not_ordered _ [] = False + not_ordered kx ((ky,_) : _) = kx >= ky + {-# INLINE not_ordered #-} + + fromList' t0 xs = foldlStrict ins t0 xs + where ins t (k,x) = insert k x t + + STRICT_1_OF_3(go) + go _ t [] = t + go _ t [(kx, x)] = insertMax kx x t + go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs + | otherwise = case create s xss of + (r, ys, []) -> go (s `shiftL` 1) (link kx x l r) ys + (r, _, ys) -> fromList' (link kx x l r) ys + + -- The create is returning a triple (tree, xs, ys). Both xs and ys + -- represent not yet processed elements and only one of them can be nonempty. + -- If ys is nonempty, the keys in ys are not ordered with respect to tree + -- and must be inserted using fromList'. Otherwise the keys have been + -- ordered so far. + STRICT_1_OF_2(create) + create _ [] = (Tip, [], []) + create s xs@(xp : xss) + | s == 1 = case xp of (kx, x) | not_ordered kx xss -> (Bin 1 kx x Tip Tip, [], xss) + | otherwise -> (Bin 1 kx x Tip Tip, xss, []) + | otherwise = case create (s `shiftR` 1) xs of + res@(_, [], _) -> res + (l, [(ky, y)], zs) -> (insertMax ky y l, [], zs) + (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys) + | otherwise -> case create (s `shiftR` 1) yss of + (r, zs, ws) -> (link ky y l r, zs, ws) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromList #-} +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] +-- > fromListWith (++) [] == empty + +fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromListWith #-} +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] +-- > fromListWithKey f [] == empty + +fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromListWithKey f xs + = foldlStrict ins empty xs + where + ins t (k,x) = insertWithKey f k x t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromListWithKey #-} +#endif + +-- | /O(n)/. Convert the map to a list of key\/value pairs. Subject to list fusion. +-- +-- > toList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] +-- > toList empty == [] + +toList :: Map k a -> [(k,a)] +toList = toAscList + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys are +-- in ascending order. Subject to list fusion. +-- +-- > toAscList (fromList [(5,"a"), (3,"b")]) == [(3,"b"), (5,"a")] + +toAscList :: Map k a -> [(k,a)] +toAscList = foldrWithKey (\k x xs -> (k,x):xs) [] + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys +-- are in descending order. Subject to list fusion. +-- +-- > toDescList (fromList [(5,"a"), (3,"b")]) == [(5,"a"), (3,"b")] + +toDescList :: Map k a -> [(k,a)] +toDescList = foldlWithKey (\xs k x -> (k,x):xs) [] + +-- List fusion for the list generating functions. +#if __GLASGOW_HASKELL__ +-- The foldrFB and foldlFB are fold{r,l}WithKey equivalents, used for list fusion. +-- They are important to convert unfused methods back, see mapFB in prelude. +foldrFB :: (k -> a -> b -> b) -> b -> Map k a -> b +foldrFB = foldrWithKey +{-# INLINE[0] foldrFB #-} +foldlFB :: (a -> k -> b -> a) -> a -> Map k b -> a +foldlFB = foldlWithKey +{-# INLINE[0] foldlFB #-} + +-- Inline assocs and toList, so that we need to fuse only toAscList. +{-# INLINE assocs #-} +{-# INLINE toList #-} + +-- The fusion is enabled up to phase 2 included. If it does not succeed, +-- convert in phase 1 the expanded elems,keys,to{Asc,Desc}List calls back to +-- elems,keys,to{Asc,Desc}List. In phase 0, we inline fold{lr}FB (which were +-- used in a list fusion, otherwise it would go away in phase 1), and let compiler +-- do whatever it wants with elems,keys,to{Asc,Desc}List -- it was forbidden to +-- inline it before phase 0, otherwise the fusion rules would not fire at all. +{-# NOINLINE[0] elems #-} +{-# NOINLINE[0] keys #-} +{-# NOINLINE[0] toAscList #-} +{-# NOINLINE[0] toDescList #-} +{-# RULES "Map.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) #-} +{-# RULES "Map.elemsBack" [1] foldrFB (\_ x xs -> x : xs) [] = elems #-} +{-# RULES "Map.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) #-} +{-# RULES "Map.keysBack" [1] foldrFB (\k _ xs -> k : xs) [] = keys #-} +{-# RULES "Map.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k,x) xs) n m) #-} +{-# RULES "Map.toAscListBack" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList #-} +{-# RULES "Map.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k,x) xs) n m) #-} +{-# RULES "Map.toDescListBack" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList #-} +#endif + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists can be done in linear time. + + Note that if [xs] is ascending that: + fromAscList xs == fromList xs + fromAscListWith f xs == fromListWith f xs +--------------------------------------------------------------------} +-- | /O(n)/. Build a map from an ascending list in linear time. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +-- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscList :: Eq k => [(k,a)] -> Map k a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromAscList #-} +#endif + +-- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] +-- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromAscListWith #-} +#endif + +-- | /O(n)/. Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False + +fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWithKey f xs + = fromDistinctAscList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromAscListWithKey #-} +#endif + + +-- | /O(n)/. Build a map from an ascending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True +-- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False + +-- For some reason, when 'singleton' is used in fromDistinctAscList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctAscList :: [(k,a)] -> Map k a +fromDistinctAscList [] = Tip +fromDistinctAscList ((kx0, x0) : xs0) = go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + STRICT_1_OF_3(go) + go _ t [] = t + go s l ((kx, x) : xs) = case create s xs of + (r, ys) -> go (s `shiftL` 1) (link kx x l r) ys + + STRICT_1_OF_2(create) + create _ [] = (Tip, []) + create s xs@(x' : xs') + | s == 1 = case x' of (kx, x) -> (Bin 1 kx x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of + (r, zs) -> (link ky y l r, zs) + + +{-------------------------------------------------------------------- + Utility functions that return sub-ranges of the original + tree. Some functions take a `Maybe value` as an argument to + allow comparisons against infinite values. These are called `blow` + (Nothing is -\infty) and `bhigh` (here Nothing is +\infty). + We use MaybeS value, which is a Maybe strict in the Just case. + + [trim blow bhigh t] A tree that is either empty or where [x > blow] + and [x < bhigh] for the value [x] of the root. + [filterGt blow t] A tree where for all values [k]. [k > blow] + [filterLt bhigh t] A tree where for all values [k]. [k < bhigh] + + [split k t] Returns two trees [l] and [r] where all keys + in [l] are <[k] and all keys in [r] are >[k]. + [splitLookup k t] Just like [split] but also returns whether [k] + was found in the tree. +--------------------------------------------------------------------} + +data MaybeS a = NothingS | JustS !a + +{-------------------------------------------------------------------- + [trim blo bhi t] trims away all subtrees that surely contain no + values between the range [blo] to [bhi]. The returned tree is either + empty or the key of the root is between @blo@ and @bhi@. +--------------------------------------------------------------------} +trim :: Ord k => MaybeS k -> MaybeS k -> Map k a -> Map k a +trim NothingS NothingS t = t +trim (JustS lk) NothingS t = greater lk t where greater lo (Bin _ k _ _ r) | k <= lo = greater lo r + greater _ t' = t' +trim NothingS (JustS hk) t = lesser hk t where lesser hi (Bin _ k _ l _) | k >= hi = lesser hi l + lesser _ t' = t' +trim (JustS lk) (JustS hk) t = middle lk hk t where middle lo hi (Bin _ k _ _ r) | k <= lo = middle lo hi r + middle lo hi (Bin _ k _ l _) | k >= hi = middle lo hi l + middle _ _ t' = t' +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE trim #-} +#endif + +-- Helper function for 'mergeWithKey'. The @'trimLookupLo' lk hk t@ performs both +-- @'trim' (JustS lk) hk t@ and @'lookup' lk t@. + +-- See Note: Type of local 'go' function +trimLookupLo :: Ord k => k -> MaybeS k -> Map k a -> (Maybe a, Map k a) +trimLookupLo lk0 mhk0 t0 = toPair $ go lk0 mhk0 t0 + where + go lk NothingS t = greater lk t + where greater :: Ord k => k -> Map k a -> StrictPair (Maybe a) (Map k a) + greater lo t'@(Bin _ kx x l r) = case compare lo kx of + LT -> lookup lo l :*: t' + EQ -> (Just x :*: r) + GT -> greater lo r + greater _ Tip = (Nothing :*: Tip) + go lk (JustS hk) t = middle lk hk t + where middle :: Ord k => k -> k -> Map k a -> StrictPair (Maybe a) (Map k a) + middle lo hi t'@(Bin _ kx x l r) = case compare lo kx of + LT | kx < hi -> lookup lo l :*: t' + | otherwise -> middle lo hi l + EQ -> Just x :*: lesser hi r + GT -> middle lo hi r + middle _ _ Tip = (Nothing :*: Tip) + + lesser :: Ord k => k -> Map k a -> Map k a + lesser hi (Bin _ k _ l _) | k >= hi = lesser hi l + lesser _ t' = t' +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE trimLookupLo #-} +#endif + + +{-------------------------------------------------------------------- + [filterGt b t] filter all keys >[b] from tree [t] + [filterLt b t] filter all keys <[b] from tree [t] +--------------------------------------------------------------------} +filterGt :: Ord k => MaybeS k -> Map k v -> Map k v +filterGt NothingS t = t +filterGt (JustS b) t = filter' b t + where filter' _ Tip = Tip + filter' b' (Bin _ kx x l r) = + case compare b' kx of LT -> link kx x (filter' b' l) r + EQ -> r + GT -> filter' b' r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE filterGt #-} +#endif + +filterLt :: Ord k => MaybeS k -> Map k v -> Map k v +filterLt NothingS t = t +filterLt (JustS b) t = filter' b t + where filter' _ Tip = Tip + filter' b' (Bin _ kx x l r) = + case compare kx b' of LT -> link kx x l (filter' b' r) + EQ -> l + GT -> filter' b' l +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE filterLt #-} +#endif + +{-------------------------------------------------------------------- + Split +--------------------------------------------------------------------} +-- | /O(log n)/. The expression (@'split' k map@) is a pair @(map1,map2)@ where +-- the keys in @map1@ are smaller than @k@ and the keys in @map2@ larger than @k@. +-- Any key equal to @k@ is found in neither @map1@ nor @map2@. +-- +-- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")]) +-- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a") +-- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") +-- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty) +-- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty) + +split :: Ord k => k -> Map k a -> (Map k a,Map k a) +split k0 t0 = k0 `seq` toPair $ go k0 t0 + where + go k t = + case t of + Tip -> (Tip :*: Tip) + Bin _ kx x l r -> case compare k kx of + LT -> let (lt :*: gt) = go k l in lt :*: link kx x gt r + GT -> let (lt :*: gt) = go k r in link kx x l lt :*: gt + EQ -> (l :*: r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE split #-} +#endif + +-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just +-- like 'split' but also returns @'lookup' k map@. +-- +-- > splitLookup 2 (fromList [(5,"a"), (3,"b")]) == (empty, Nothing, fromList [(3,"b"), (5,"a")]) +-- > splitLookup 3 (fromList [(5,"a"), (3,"b")]) == (empty, Just "b", singleton 5 "a") +-- > splitLookup 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Nothing, singleton 5 "a") +-- > splitLookup 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", Just "a", empty) +-- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty) + +splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a) +splitLookup k t = k `seq` + case t of + Tip -> (Tip,Nothing,Tip) + Bin _ kx x l r -> case compare k kx of + LT -> let (lt,z,gt) = splitLookup k l + gt' = link kx x gt r + in gt' `seq` (lt,z,gt') + GT -> let (lt,z,gt) = splitLookup k r + lt' = link kx x l lt + in lt' `seq` (lt',z,gt) + EQ -> (l,Just x,r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE splitLookup #-} +#endif + +{-------------------------------------------------------------------- + Utility functions that maintain the balance properties of the tree. + All constructors assume that all values in [l] < [k] and all values + in [r] > [k], and that [l] and [r] are valid trees. + + In order of sophistication: + [Bin sz k x l r] The type constructor. + [bin k x l r] Maintains the correct size, assumes that both [l] + and [r] are balanced with respect to each other. + [balance k x l r] Restores the balance and size. + Assumes that the original tree was balanced and + that [l] or [r] has changed by at most one element. + [link k x l r] Restores balance and size. + + Furthermore, we can construct a new tree from two trees. Both operations + assume that all values in [l] < all values in [r] and that [l] and [r] + are valid: + [glue l r] Glues [l] and [r] together. Assumes that [l] and + [r] are already balanced with respect to each other. + [merge l r] Merges two trees and restores balance. + + Note: in contrast to Adam's paper, we use (<=) comparisons instead + of (<) comparisons in [link], [merge] and [balance]. + Quickcheck (on [difference]) showed that this was necessary in order + to maintain the invariants. It is quite unsatisfactory that I haven't + been able to find out why this is actually the case! Fortunately, it + doesn't hurt to be a bit more conservative. +--------------------------------------------------------------------} + +{-------------------------------------------------------------------- + Link +--------------------------------------------------------------------} +link :: k -> a -> Map k a -> Map k a -> Map k a +link kx x Tip r = insertMin kx x r +link kx x l Tip = insertMax kx x l +link kx x l@(Bin sizeL ky y ly ry) r@(Bin sizeR kz z lz rz) + | delta*sizeL < sizeR = balanceL kz z (link kx x l lz) rz + | delta*sizeR < sizeL = balanceR ky y ly (link kx x ry r) + | otherwise = bin kx x l r + + +-- insertMin and insertMax don't perform potentially expensive comparisons. +insertMax,insertMin :: k -> a -> Map k a -> Map k a +insertMax kx x t + = case t of + Tip -> singleton kx x + Bin _ ky y l r + -> balanceR ky y l (insertMax kx x r) + +insertMin kx x t + = case t of + Tip -> singleton kx x + Bin _ ky y l r + -> balanceL ky y (insertMin kx x l) r + +{-------------------------------------------------------------------- + [merge l r]: merges two trees. +--------------------------------------------------------------------} +merge :: Map k a -> Map k a -> Map k a +merge Tip r = r +merge l Tip = l +merge l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) + | delta*sizeL < sizeR = balanceL ky y (merge l ly) ry + | delta*sizeR < sizeL = balanceR kx x lx (merge rx r) + | otherwise = glue l r + +{-------------------------------------------------------------------- + [glue l r]: glues two trees together. + Assumes that [l] and [r] are already balanced with respect to each other. +--------------------------------------------------------------------} +glue :: Map k a -> Map k a -> Map k a +glue Tip r = r +glue l Tip = l +glue l r + | size l > size r = let ((km,m),l') = deleteFindMax l in balanceR km m l' r + | otherwise = let ((km,m),r') = deleteFindMin r in balanceL km m l r' + + +-- | /O(log n)/. Delete and find the minimal element. +-- +-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")]) +-- > deleteFindMin Error: can not return the minimal element of an empty map + +deleteFindMin :: Map k a -> ((k,a),Map k a) +deleteFindMin t + = case t of + Bin _ k x Tip r -> ((k,x),r) + Bin _ k x l r -> let (km,l') = deleteFindMin l in (km,balanceR k x l' r) + Tip -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) + +-- | /O(log n)/. Delete and find the maximal element. +-- +-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")]) +-- > deleteFindMax empty Error: can not return the maximal element of an empty map + +deleteFindMax :: Map k a -> ((k,a),Map k a) +deleteFindMax t + = case t of + Bin _ k x l Tip -> ((k,x),l) + Bin _ k x l r -> let (km,r') = deleteFindMax r in (km,balanceL k x l r') + Tip -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) + + +{-------------------------------------------------------------------- + [balance l x r] balances two trees with value x. + The sizes of the trees should balance after decreasing the + size of one of them. (a rotation). + + [delta] is the maximal relative difference between the sizes of + two trees, it corresponds with the [w] in Adams' paper. + [ratio] is the ratio between an outer and inner sibling of the + heavier subtree in an unbalanced setting. It determines + whether a double or single rotation should be performed + to restore balance. It is corresponds with the inverse + of $\alpha$ in Adam's article. + + Note that according to the Adam's paper: + - [delta] should be larger than 4.646 with a [ratio] of 2. + - [delta] should be larger than 3.745 with a [ratio] of 1.534. + + But the Adam's paper is erroneous: + - It can be proved that for delta=2 and delta>=5 there does + not exist any ratio that would work. + - Delta=4.5 and ratio=2 does not work. + + That leaves two reasonable variants, delta=3 and delta=4, + both with ratio=2. + + - A lower [delta] leads to a more 'perfectly' balanced tree. + - A higher [delta] performs less rebalancing. + + In the benchmarks, delta=3 is faster on insert operations, + and delta=4 has slightly better deletes. As the insert speedup + is larger, we currently use delta=3. + +--------------------------------------------------------------------} +delta,ratio :: Int +delta = 3 +ratio = 2 + +-- The balance function is equivalent to the following: +-- +-- balance :: k -> a -> Map k a -> Map k a -> Map k a +-- balance k x l r +-- | sizeL + sizeR <= 1 = Bin sizeX k x l r +-- | sizeR > delta*sizeL = rotateL k x l r +-- | sizeL > delta*sizeR = rotateR k x l r +-- | otherwise = Bin sizeX k x l r +-- where +-- sizeL = size l +-- sizeR = size r +-- sizeX = sizeL + sizeR + 1 +-- +-- rotateL :: a -> b -> Map a b -> Map a b -> Map a b +-- rotateL k x l r@(Bin _ _ _ ly ry) | size ly < ratio*size ry = singleL k x l r +-- | otherwise = doubleL k x l r +-- +-- rotateR :: a -> b -> Map a b -> Map a b -> Map a b +-- rotateR k x l@(Bin _ _ _ ly ry) r | size ry < ratio*size ly = singleR k x l r +-- | otherwise = doubleR k x l r +-- +-- singleL, singleR :: a -> b -> Map a b -> Map a b -> Map a b +-- singleL k1 x1 t1 (Bin _ k2 x2 t2 t3) = bin k2 x2 (bin k1 x1 t1 t2) t3 +-- singleR k1 x1 (Bin _ k2 x2 t1 t2) t3 = bin k2 x2 t1 (bin k1 x1 t2 t3) +-- +-- doubleL, doubleR :: a -> b -> Map a b -> Map a b -> Map a b +-- doubleL k1 x1 t1 (Bin _ k2 x2 (Bin _ k3 x3 t2 t3) t4) = bin k3 x3 (bin k1 x1 t1 t2) (bin k2 x2 t3 t4) +-- doubleR k1 x1 (Bin _ k2 x2 t1 (Bin _ k3 x3 t2 t3)) t4 = bin k3 x3 (bin k2 x2 t1 t2) (bin k1 x1 t3 t4) +-- +-- It is only written in such a way that every node is pattern-matched only once. + +balance :: k -> a -> Map k a -> Map k a -> Map k a +balance k x l r = case l of + Tip -> case r of + Tip -> Bin 1 k x Tip Tip + (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r + (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr + (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) + (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) + | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr + | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + + (Bin ls lk lx ll lr) -> case r of + Tip -> case (ll, lr) of + (Tip, Tip) -> Bin 2 k x l Tip + (Tip, (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) + ((Bin _ _ _ _ _), Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) + ((Bin lls _ _ _ _), (Bin lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) + | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) + (Bin rs rk rx rl rr) + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr + | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + (_, _) -> error "Failure in Data.Map.balance" + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) + | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) + (_, _) -> error "Failure in Data.Map.balance" + | otherwise -> Bin (1+ls+rs) k x l r +{-# NOINLINE balance #-} + +-- Functions balanceL and balanceR are specialised versions of balance. +-- balanceL only checks whether the left subtree is too big, +-- balanceR only checks whether the right subtree is too big. + +-- balanceL is called when left subtree might have been inserted to or when +-- right subtree might have been deleted from. +balanceL :: k -> a -> Map k a -> Map k a -> Map k a +balanceL k x l r = case r of + Tip -> case l of + Tip -> Bin 1 k x Tip Tip + (Bin _ _ _ Tip Tip) -> Bin 2 k x l Tip + (Bin _ lk lx Tip (Bin _ lrk lrx _ _)) -> Bin 3 lrk lrx (Bin 1 lk lx Tip Tip) (Bin 1 k x Tip Tip) + (Bin _ lk lx ll@(Bin _ _ _ _ _) Tip) -> Bin 3 lk lx ll (Bin 1 k x Tip Tip) + (Bin ls lk lx ll@(Bin lls _ _ _ _) lr@(Bin lrs lrk lrx lrl lrr)) + | lrs < ratio*lls -> Bin (1+ls) lk lx ll (Bin (1+lrs) k x lr Tip) + | otherwise -> Bin (1+ls) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+size lrr) k x lrr Tip) + + (Bin rs _ _ _ _) -> case l of + Tip -> Bin (1+rs) k x Tip r + + (Bin ls lk lx ll lr) + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _ _, Bin lrs lrk lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lk lx ll (Bin (1+rs+lrs) k x lr r) + | otherwise -> Bin (1+ls+rs) lrk lrx (Bin (1+lls+size lrl) lk lx ll lrl) (Bin (1+rs+size lrr) k x lrr r) + (_, _) -> error "Failure in Data.Map.balanceL" + | otherwise -> Bin (1+ls+rs) k x l r +{-# NOINLINE balanceL #-} + +-- balanceR is called when right subtree might have been inserted to or when +-- left subtree might have been deleted from. +balanceR :: k -> a -> Map k a -> Map k a -> Map k a +balanceR k x l r = case l of + Tip -> case r of + Tip -> Bin 1 k x Tip Tip + (Bin _ _ _ Tip Tip) -> Bin 2 k x Tip r + (Bin _ rk rx Tip rr@(Bin _ _ _ _ _)) -> Bin 3 rk rx (Bin 1 k x Tip Tip) rr + (Bin _ rk rx (Bin _ rlk rlx _ _) Tip) -> Bin 3 rlk rlx (Bin 1 k x Tip Tip) (Bin 1 rk rx Tip Tip) + (Bin rs rk rx rl@(Bin rls rlk rlx rll rlr) rr@(Bin rrs _ _ _ _)) + | rls < ratio*rrs -> Bin (1+rs) rk rx (Bin (1+rls) k x Tip rl) rr + | otherwise -> Bin (1+rs) rlk rlx (Bin (1+size rll) k x Tip rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + + (Bin ls _ _ _ _) -> case r of + Tip -> Bin (1+ls) k x l Tip + + (Bin rs rk rx rl rr) + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlk rlx rll rlr, Bin rrs _ _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rk rx (Bin (1+ls+rls) k x l rl) rr + | otherwise -> Bin (1+ls+rs) rlk rlx (Bin (1+ls+size rll) k x l rll) (Bin (1+rrs+size rlr) rk rx rlr rr) + (_, _) -> error "Failure in Data.Map.balanceR" + | otherwise -> Bin (1+ls+rs) k x l r +{-# NOINLINE balanceR #-} + + +{-------------------------------------------------------------------- + The bin constructor maintains the size of the tree +--------------------------------------------------------------------} +bin :: k -> a -> Map k a -> Map k a -> Map k a +bin k x l r + = Bin (size l + size r + 1) k x l r +{-# INLINE bin #-} + + +{-------------------------------------------------------------------- + Eq converts the tree to a list. In a lazy setting, this + actually seems one of the faster methods to compare two trees + and it is certainly the simplest :-) +--------------------------------------------------------------------} +instance (Eq k,Eq a) => Eq (Map k a) where + t1 == t2 = (size t1 == size t2) && (toAscList t1 == toAscList t2) + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance (Ord k, Ord v) => Ord (Map k v) where + compare m1 m2 = compare (toAscList m1) (toAscList m2) + +{-------------------------------------------------------------------- + Functor +--------------------------------------------------------------------} +instance Functor (Map k) where + fmap f m = map f m + +instance Traversable (Map k) where + traverse f = traverseWithKey (\_ -> f) + {-# INLINE traverse #-} + +instance Foldable.Foldable (Map k) where + fold t = go t + where go Tip = mempty + go (Bin 1 _ v _ _) = v + go (Bin _ _ v l r) = go l `mappend` (v `mappend` go r) + {-# INLINABLE fold #-} + foldr = foldr + {-# INLINE foldr #-} + foldl = foldl + {-# INLINE foldl #-} + foldMap f t = go t + where go Tip = mempty + go (Bin 1 _ v _ _) = f v + go (Bin _ _ v l r) = go l `mappend` (f v `mappend` go r) + {-# INLINE foldMap #-} + +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 #-} + + +-- | /O(1)/. Decompose a map into pieces based on the structure of the underlying +-- tree. This function is useful for consuming a map in parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that the pieces +-- returned will be in ascending order (all elements in the first submap less than all +-- elements in the second, and so on). +-- +-- Examples: +-- +-- > splitRoot (fromList (zip [1..6] ['a'..])) == +-- > [fromList [(1,'a'),(2,'b'),(3,'c')],fromList [(4,'d')],fromList [(5,'e'),(6,'f')]] +-- +-- > splitRoot empty == [] +-- +-- Note that the current implementation does not return more than three submaps, +-- but you should not depend on this behaviour because it can change in the +-- future without notice. +splitRoot :: Map k b -> [Map k b] +splitRoot orig = + case orig of + Tip -> [] + Bin _ k v l r -> [l, singleton k v, r] +{-# INLINE splitRoot #-} 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,230 @@ +{-# 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 + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , M.filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + +#if defined(TESTING) + -- * Internals + , bin + , balanced + , link + , merge +#endif + + ) where + +import 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,1182 @@ +{-# 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 + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + , fromSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Filter + , filter + , filterWithKey + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , deleteAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + , valid + +#if defined(TESTING) + -- * Internals + , bin + , balanced + , link + , merge +#endif + ) where + +import Prelude hiding (lookup,map,filter,foldr,foldl,null) + +import Data.Map.Base hiding + ( findWithDefault + , singleton + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , unionWith + , unionWithKey + , unionsWith + , differenceWith + , differenceWithKey + , intersectionWith + , intersectionWithKey + , mergeWithKey + , map + , mapWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeysWith + , fromSet + , fromList + , fromListWith + , fromListWithKey + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + , updateAt + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + ) +import qualified Data.Set.Base as Set +import Data.StrictPair +import Data.Bits (shiftL, shiftR) + +-- 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_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_2_OF_3(fn) fn arg1 arg2 _ | arg1 `seq` arg2 `seq` False = undefined +#define STRICT_2_OF_4(fn) fn _ arg _ _ | arg `seq` False = undefined + +-- $strictness +-- +-- This module satisfies the following strictness properties: +-- +-- 1. Key arguments are evaluated to WHNF; +-- +-- 2. Keys and values are evaluated to WHNF before they are stored in +-- the map. +-- +-- Here's an example illustrating the first property: +-- +-- > delete undefined m == undefined +-- +-- Here are some examples that illustrate the second property: +-- +-- > map (\ v -> undefined) m == undefined -- m is not empty +-- > mapKeys (\ k -> undefined) m == undefined -- m is not empty + +{-------------------------------------------------------------------- + Query +--------------------------------------------------------------------} + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' +-- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' + +-- See Map.Base.Note: Local 'go' functions and capturing +findWithDefault :: Ord k => a -> k -> Map k a -> a +findWithDefault def k = k `seq` go + where + go Tip = def + go (Bin _ kx x l r) = case compare k kx of + LT -> go l + GT -> go r + EQ -> x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE findWithDefault #-} +#else +{-# INLINE findWithDefault #-} +#endif + +{-------------------------------------------------------------------- + Construction +--------------------------------------------------------------------} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList [(1, 'a')] +-- > size (singleton 1 'a') == 1 + +singleton :: k -> a -> Map k a +singleton k x = x `seq` Bin 1 k x Tip Tip +{-# INLINE singleton #-} + +{-------------------------------------------------------------------- + Insertion +--------------------------------------------------------------------} +-- | /O(log n)/. Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value. 'insert' is equivalent to +-- @'insertWith' 'const'@. +-- +-- > insert 5 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'x')] +-- > insert 7 'x' (fromList [(5,'a'), (3,'b')]) == fromList [(3, 'b'), (5, 'a'), (7, 'x')] +-- > insert 5 'x' empty == singleton 5 'x' + +-- See Map.Base.Note: Type of local 'go' function +insert :: Ord k => k -> a -> Map k a -> Map k a +insert = go + where + go :: Ord k => k -> a -> Map k a -> Map k a + STRICT_1_2_OF_3(go) + go kx x Tip = singleton kx x + go kx x (Bin sz ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go kx x l) r + GT -> balanceR ky y l (go kx x r) + EQ -> Bin sz kx x l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insert #-} +#else +{-# INLINE insert #-} +#endif + +-- | /O(log n)/. Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key, f new_value old_value)@. +-- +-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] +-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" + +insertWith :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWith f = insertWithKey (\_ x' y' -> f x' y') +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertWith #-} +#else +{-# INLINE insertWith #-} +#endif + +-- | /O(log n)/. Insert with a function, combining key, new value and old value. +-- @'insertWithKey' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert the pair @(key,f key new_value old_value)@. +-- Note that the key passed to f is the same key passed to 'insertWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] +-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" + +-- See Map.Base.Note: Type of local 'go' function +insertWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a +insertWithKey = go + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a + STRICT_2_OF_4(go) + go _ kx x Tip = singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> balanceL ky y (go f kx x l) r + GT -> balanceR ky y l (go f kx x r) + EQ -> let x' = f kx x y + in x' `seq` Bin sy kx x' l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertWithKey #-} +#else +{-# INLINE insertWithKey #-} +#endif + +-- | /O(log n)/. Combines insert operation with old value retrieval. +-- The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- and the second element equal to (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) +-- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) +-- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) + +-- See Map.Base.Note: Type of local 'go' function +insertLookupWithKey :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a + -> (Maybe a, Map k a) +insertLookupWithKey f0 kx0 x0 t0 = toPair $ go f0 kx0 x0 t0 + where + go :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> StrictPair (Maybe a) (Map k a) + STRICT_2_OF_4(go) + go _ kx x Tip = Nothing :*: singleton kx x + go f kx x (Bin sy ky y l r) = + case compare kx ky of + LT -> let (found :*: l') = go f kx x l + in found :*: balanceL ky y l' r + GT -> let (found :*: r') = go f kx x r + in found :*: balanceR ky y l r' + EQ -> let x' = f kx x y + in x' `seq` (Just y :*: Bin sy kx x' l r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE insertLookupWithKey #-} +#else +{-# INLINE insertLookupWithKey #-} +#endif + +{-------------------------------------------------------------------- + Deletion +--------------------------------------------------------------------} + +-- | /O(log n)/. Update a value at a specific key with the result of the provided function. +-- When the key is not +-- a member of the map, the original map is returned. +-- +-- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > adjust ("new " ++) 7 empty == empty + +adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a +adjust f = adjustWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE adjust #-} +#else +{-# INLINE adjust #-} +#endif + +-- | /O(log n)/. Adjust a value at a specific key. When the key is not +-- a member of the map, the original map is returned. +-- +-- > let f key x = (show key) ++ ":new " ++ x +-- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > adjustWithKey f 7 empty == empty + +adjustWithKey :: Ord k => (k -> a -> a) -> k -> Map k a -> Map k a +adjustWithKey f = updateWithKey (\k' x' -> Just (f k' x')) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE adjustWithKey #-} +#else +{-# INLINE adjustWithKey #-} +#endif + +-- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] +-- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a +update f = updateWithKey (\_ x -> f x) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE update #-} +#else +{-# INLINE update #-} +#endif + +-- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the +-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', +-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound +-- to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +-- See Map.Base.Note: Type of local 'go' function +updateWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a +updateWithKey = go + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> Map k a + STRICT_2_OF_3(go) + go _ _ Tip = Tip + go f k(Bin sx kx x l r) = + case compare k kx of + LT -> balanceR kx x (go f k l) r + GT -> balanceL kx x l (go f k r) + EQ -> case f kx x of + Just x' -> x' `seq` Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE updateWithKey #-} +#else +{-# INLINE updateWithKey #-} +#endif + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "5:new a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") + +-- See Map.Base.Note: Type of local 'go' function +updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) +updateLookupWithKey f0 k0 t0 = toPair $ go f0 k0 t0 + where + go :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> StrictPair (Maybe a) (Map k a) + STRICT_2_OF_3(go) + go _ _ Tip = (Nothing :*: Tip) + go f k (Bin sx kx x l r) = + case compare k kx of + LT -> let (found :*: l') = go f k l + in found :*: balanceR kx x l' r + GT -> let (found :*: r') = go f k r + in found :*: balanceL kx x l r' + EQ -> case f kx x of + Just x' -> x' `seq` (Just x' :*: Bin sx kx x' l r) + Nothing -> (Just x :*: glue l r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE updateLookupWithKey #-} +#else +{-# INLINE updateLookupWithKey #-} +#endif + +-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. +-- 'alter' can be used to insert, delete, or update a value in a 'Map'. +-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "c")] + +-- See Map.Base.Note: Type of local 'go' function +alter :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a +alter = go + where + go :: Ord k => (Maybe a -> Maybe a) -> k -> Map k a -> Map k a + STRICT_2_OF_3(go) + go f k Tip = case f Nothing of + Nothing -> Tip + Just x -> singleton k x + + go f k (Bin sx kx x l r) = case compare k kx of + LT -> balance kx x (go f k l) r + GT -> balance kx x l (go f k r) + EQ -> case f (Just x) of + Just x' -> x' `seq` Bin sx kx x' l r + Nothing -> glue l r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE alter #-} +#else +{-# INLINE alter #-} +#endif + +{-------------------------------------------------------------------- + Indexing +--------------------------------------------------------------------} + +-- | /O(log n)/. Update the element at /index/. Calls 'error' when an +-- invalid index is used. +-- +-- > updateAt (\ _ _ -> Just "x") 0 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "x"), (5, "a")] +-- > updateAt (\ _ _ -> Just "x") 1 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "x")] +-- > updateAt (\ _ _ -> Just "x") 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > updateAt (\ _ _ -> Just "x") (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > updateAt (\_ _ -> Nothing) 0 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +-- > updateAt (\_ _ -> Nothing) 1 (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" +-- > updateAt (\_ _ -> Nothing) 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range +-- > updateAt (\_ _ -> Nothing) (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range + +updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a +updateAt f i t = i `seq` + case t of + Tip -> error "Map.updateAt: index out of range" + Bin sx kx x l r -> case compare i sizeL of + LT -> balanceR kx x (updateAt f i l) r + GT -> balanceL kx x l (updateAt f (i-sizeL-1) r) + EQ -> case f kx x of + Just x' -> x' `seq` Bin sx kx x' l r + Nothing -> glue l r + where + sizeL = size l + +{-------------------------------------------------------------------- + Minimal, Maximal +--------------------------------------------------------------------} + +-- | /O(log n)/. Update the value at the minimal key. +-- +-- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")] +-- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +updateMin :: (a -> Maybe a) -> Map k a -> Map k a +updateMin f m + = updateMinWithKey (\_ x -> f x) m + +-- | /O(log n)/. Update the value at the maximal key. +-- +-- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")] +-- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" + +updateMax :: (a -> Maybe a) -> Map k a -> Map k a +updateMax f m + = updateMaxWithKey (\_ x -> f x) m + + +-- | /O(log n)/. Update the value at the minimal key. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" + +updateMinWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMinWithKey _ Tip = Tip +updateMinWithKey f (Bin sx kx x Tip r) = case f kx x of + Nothing -> r + Just x' -> x' `seq` Bin sx kx x' Tip r +updateMinWithKey f (Bin _ kx x l r) = balanceR kx x (updateMinWithKey f l) r + +-- | /O(log n)/. Update the value at the maximal key. +-- +-- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")] +-- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" + +updateMaxWithKey :: (k -> a -> Maybe a) -> Map k a -> Map k a +updateMaxWithKey _ Tip = Tip +updateMaxWithKey f (Bin sx kx x l Tip) = case f kx x of + Nothing -> l + Just x' -> x' `seq` Bin sx kx x' l Tip +updateMaxWithKey f (Bin _ kx x l r) = balanceL kx x l (updateMaxWithKey f r) + +{-------------------------------------------------------------------- + Union. +--------------------------------------------------------------------} + +-- | The union of a list of maps, with a combining operation: +-- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). +-- +-- > unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] +-- > == fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] + +unionsWith :: Ord k => (a->a->a) -> [Map k a] -> Map k a +unionsWith f ts + = foldlStrict (unionWith f) empty ts +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unionsWith #-} +#endif + +{-------------------------------------------------------------------- + Union with a combining function +--------------------------------------------------------------------} +-- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] + +unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWith f m1 m2 + = unionWithKey (\_ x y -> f x y) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unionWith #-} +#endif + +-- | /O(n+m)/. +-- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] + +unionWithKey :: Ord k => (k -> a -> a -> a) -> Map k a -> Map k a -> Map k a +unionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) id id t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE unionWithKey #-} +#endif + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} + +-- | /O(n+m)/. Difference with a combining function. +-- When two equal keys are +-- encountered, the combining function is applied to the values of these keys. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) +-- > == singleton 3 "b:B" + +differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWith f m1 m2 + = differenceWithKey (\_ x y -> f x y) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE differenceWith #-} +#endif + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) +-- > == singleton 3 "3:b|B" + +differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a +differenceWithKey f t1 t2 = mergeWithKey f id (const Tip) t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE differenceWithKey #-} +#endif + + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} + +-- | /O(n+m)/. Intersection with a combining function. The implementation uses +-- an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" + +intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWith f m1 m2 + = intersectionWithKey (\_ x y -> f x y) m1 m2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE intersectionWith #-} +#endif + +-- | /O(n+m)/. Intersection with a combining function. The implementation uses +-- an efficient /hedge/ algorithm comparable with /hedge-union/. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" + + +intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c +intersectionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) (const Tip) (const Tip) t1 t2 +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE intersectionWithKey #-} +#endif + + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | /O(n+m)/. A high-performance universal combining function. This function +-- is used to define 'unionWith', 'unionWithKey', 'differenceWith', +-- 'differenceWithKey', 'intersectionWith', 'intersectionWithKey' and can be +-- used to define other custom combine functions. +-- +-- Please make sure you know what is going on when using 'mergeWithKey', +-- otherwise you can be surprised by unexpected code growth or even +-- corruption of the data structure. +-- +-- When 'mergeWithKey' is given three arguments, it is inlined to the call +-- site. You should therefore use 'mergeWithKey' only to define your custom +-- combining functions. For example, you could define 'unionWithKey', +-- 'differenceWithKey' and 'intersectionWithKey' as +-- +-- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 +-- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 +-- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 +-- +-- When calling @'mergeWithKey' combine only1 only2@, a function combining two +-- 'IntMap's is created, such that +-- +-- * if a key is present in both maps, it is passed with both corresponding +-- values to the @combine@ function. Depending on the result, the key is either +-- present in the result with specified value, or is left out; +-- +-- * a nonempty subtree present only in the first map is passed to @only1@ and +-- the output is added to the result; +-- +-- * a nonempty subtree present only in the second map is passed to @only2@ and +-- the output is added to the result. +-- +-- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. +-- The values can be modified arbitrarily. Most common variants of @only1@ and +-- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or +-- @'filterWithKey' f@ could be used for any @f@. + +mergeWithKey :: Ord k => (k -> a -> b -> Maybe c) -> (Map k a -> Map k c) -> (Map k b -> Map k c) + -> Map k a -> Map k b -> Map k c +mergeWithKey f g1 g2 = go + where + go Tip t2 = g2 t2 + go t1 Tip = g1 t1 + go t1 t2 = hedgeMerge NothingS NothingS t1 t2 + + hedgeMerge _ _ t1 Tip = g1 t1 + hedgeMerge blo bhi Tip (Bin _ kx x l r) = g2 $ link kx x (filterGt blo l) (filterLt bhi r) + hedgeMerge blo bhi (Bin _ kx x l r) t2 = let l' = hedgeMerge blo bmi l (trim blo bmi t2) + (found, trim_t2) = trimLookupLo kx bhi t2 + r' = hedgeMerge bmi bhi r trim_t2 + in case found of + Nothing -> case g1 (singleton kx x) of + Tip -> merge l' r' + (Bin _ _ x' Tip Tip) -> link kx x' l' r' + _ -> error "mergeWithKey: Given function only1 does not fulfil required conditions (see documentation)" + Just x2 -> case f kx x x2 of + Nothing -> merge l' r' + Just x' -> x' `seq` link kx x' l' r' + where bmi = JustS kx +{-# INLINE mergeWithKey #-} + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} + +-- | /O(n)/. Map values and collect the 'Just' results. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a" + +mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b +mapMaybe f = mapMaybeWithKey (\_ x -> f x) + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +-- +-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing +-- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" + +mapMaybeWithKey :: (k -> a -> Maybe b) -> Map k a -> Map k b +mapMaybeWithKey _ Tip = Tip +mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of + Just y -> y `seq` link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r) + Nothing -> merge (mapMaybeWithKey f l) (mapMaybeWithKey f r) + +-- | /O(n)/. Map values and separate the 'Left' and 'Right' results. +-- +-- > let f a = if a < "c" then Left a else Right a +-- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")]) +-- > +-- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) + +mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEither f m + = mapEitherWithKey (\_ x -> f x) m + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- +-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) +-- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")]) +-- > +-- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) +-- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) + +mapEitherWithKey :: (k -> a -> Either b c) -> Map k a -> (Map k b, Map k c) +mapEitherWithKey f0 t0 = toPair $ go f0 t0 + where + go _ Tip = (Tip :*: Tip) + go f (Bin _ kx x l r) = case f kx x of + Left y -> y `seq` (link kx y l1 r1 :*: merge l2 r2) + Right z -> z `seq` (merge l1 r1 :*: link kx z l2 r2) + where + (l1 :*: l2) = go f l + (r1 :*: r2) = go f r + +{-------------------------------------------------------------------- + Mapping +--------------------------------------------------------------------} +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] + +map :: (a -> b) -> Map k a -> Map k b +map _ Tip = Tip +map f (Bin sx kx x l r) = let x' = f x in x' `seq` Bin sx kx x' (map f l) (map f r) + +-- | /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. +-- +-- If the keys of the list are ordered, linear-time implementation is used, +-- with the performance equal to 'fromDistinctAscList'. +-- +-- > fromList [] == empty +-- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")] +-- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")] + +-- For some reason, when 'singleton' is used in fromList or in +-- create, it is not inlined, so we inline it manually. +fromList :: Ord k => [(k,a)] -> Map k a +fromList [] = Tip +fromList [(kx, x)] = x `seq` Bin 1 kx x Tip Tip +fromList ((kx0, x0) : xs0) | not_ordered kx0 xs0 = x0 `seq` fromList' (Bin 1 kx0 x0 Tip Tip) xs0 + | otherwise = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + not_ordered _ [] = False + not_ordered kx ((ky,_) : _) = kx >= ky + {-# INLINE not_ordered #-} + + fromList' t0 xs = foldlStrict ins t0 xs + where ins t (k,x) = insert k x t + + STRICT_1_OF_3(go) + go _ t [] = t + go _ t [(kx, x)] = x `seq` insertMax kx x t + go s l xs@((kx, x) : xss) | not_ordered kx xss = fromList' l xs + | otherwise = case create s xss of + (r, ys, []) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys + (r, _, ys) -> x `seq` fromList' (link kx x l r) ys + + -- The create is returning a triple (tree, xs, ys). Both xs and ys + -- represent not yet processed elements and only one of them can be nonempty. + -- If ys is nonempty, the keys in ys are not ordered with respect to tree + -- and must be inserted using fromList'. Otherwise the keys have been + -- ordered so far. + STRICT_1_OF_2(create) + create _ [] = (Tip, [], []) + create s xs@(xp : xss) + | s == 1 = case xp of (kx, x) | not_ordered kx xss -> x `seq` (Bin 1 kx x Tip Tip, [], xss) + | otherwise -> x `seq` (Bin 1 kx x Tip Tip, xss, []) + | otherwise = case create (s `shiftR` 1) xs of + res@(_, [], _) -> res + (l, [(ky, y)], zs) -> y `seq` (insertMax ky y l, [], zs) + (l, ys@((ky, y):yss), _) | not_ordered ky yss -> (l, [], ys) + | otherwise -> case create (s `shiftR` 1) yss of + (r, zs, ws) -> y `seq` (link ky y l r, zs, ws) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromList #-} +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] +-- > fromListWith (++) [] == empty + +fromListWith :: Ord k => (a -> a -> a) -> [(k,a)] -> Map k a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromListWith #-} +#endif + +-- | /O(n*log n)/. Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "3ab"), (5, "5a5ba")] +-- > fromListWithKey f [] == empty + +fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromListWithKey f xs + = foldlStrict ins empty xs + where + ins t (k,x) = insertWithKey f k x t +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromListWithKey #-} +#endif + +{-------------------------------------------------------------------- + Building trees from ascending/descending lists can be done in linear time. + + Note that if [xs] is ascending that: + fromAscList xs == fromList xs + fromAscListWith f xs == fromListWith f xs +--------------------------------------------------------------------} +-- | /O(n)/. Build a map from an ascending list in linear time. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +-- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscList :: Eq k => [(k,a)] -> Map k a +fromAscList xs + = fromAscListWithKey (\_ x _ -> x) xs +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromAscList #-} +#endif + +-- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] +-- > valid (fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")]) == True +-- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False + +fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWith f xs + = fromAscListWithKey (\_ x y -> f x y) xs +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromAscListWith #-} +#endif + +-- | /O(n)/. Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] == fromList [(3, "b"), (5, "5:b5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False + +fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a +fromAscListWithKey f xs + = fromDistinctAscList (combineEq f xs) + where + -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs] + combineEq _ xs' + = case xs' of + [] -> [] + [x] -> [x] + (x:xx) -> combineEq' x xx + + combineEq' z [] = [z] + combineEq' z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs' + | otherwise = z:combineEq' x xs' +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE fromAscListWithKey #-} +#endif + +-- | /O(n)/. Build a map from an ascending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > valid (fromDistinctAscList [(3,"b"), (5,"a")]) == True +-- > valid (fromDistinctAscList [(3,"b"), (5,"a"), (5,"b")]) == False + +-- For some reason, when 'singleton' is used in fromDistinctAscList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctAscList :: [(k,a)] -> Map k a +fromDistinctAscList [] = Tip +fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0 + where + STRICT_1_OF_3(go) + go _ t [] = t + go s l ((kx, x) : xs) = case create s xs of + (r, ys) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys + + STRICT_1_OF_2(create) + create _ [] = (Tip, []) + create s xs@(x' : xs') + | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of + (r, zs) -> y `seq` (link ky y l r, zs) diff --git a/Data/Sequence.hs b/Data/Sequence.hs new file mode 100644 --- /dev/null +++ b/Data/Sequence.hs @@ -0,0 +1,1804 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +----------------------------------------------------------------------------- +-- | +-- Module : Data.Sequence +-- Copyright : (c) Ross Paterson 2005 +-- (c) Louis Wasserman 2009 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : experimental +-- Portability : portable +-- +-- General purpose finite sequences. +-- Apart from being finite and having strict operations, sequences +-- also differ from lists in supporting a wider variety of operations +-- efficiently. +-- +-- An amortized running time is given for each operation, with /n/ referring +-- to the length of the sequence and /i/ being the integral index used by +-- some operations. These bounds hold even in a persistent (shared) setting. +-- +-- The implementation uses 2-3 finger trees annotated with sizes, +-- as described in section 4.2 of +-- +-- * Ralf Hinze and Ross Paterson, +-- \"Finger trees: a simple general-purpose data structure\", +-- /Journal of Functional Programming/ 16:2 (2006) pp 197-217. +-- +-- +-- /Note/: Many of these operations have the same names as similar +-- operations on lists in the "Prelude". The ambiguity may be resolved +-- using either qualification or the @hiding@ clause. +-- +----------------------------------------------------------------------------- + +module Data.Sequence ( +#if !defined(TESTING) + Seq, +#else + Seq(..), Elem(..), FingerTree(..), Node(..), Digit(..), +#endif + -- * Construction + empty, -- :: Seq a + singleton, -- :: a -> Seq a + (<|), -- :: a -> Seq a -> Seq a + (|>), -- :: Seq a -> a -> Seq a + (><), -- :: Seq a -> Seq a -> Seq a + fromList, -- :: [a] -> Seq a + -- ** Repetition + replicate, -- :: Int -> a -> Seq a + replicateA, -- :: Applicative f => Int -> f a -> f (Seq a) + replicateM, -- :: Monad m => Int -> m a -> m (Seq a) + -- ** Iterative construction + iterateN, -- :: Int -> (a -> a) -> a -> Seq a + unfoldr, -- :: (b -> Maybe (a, b)) -> b -> Seq a + unfoldl, -- :: (b -> Maybe (b, a)) -> b -> Seq a + -- * Deconstruction + -- | Additional functions for deconstructing sequences are available + -- via the 'Foldable' instance of 'Seq'. + + -- ** Queries + null, -- :: Seq a -> Bool + length, -- :: Seq a -> Int + -- ** Views + ViewL(..), + viewl, -- :: Seq a -> ViewL a + ViewR(..), + viewr, -- :: Seq a -> ViewR a + -- * Scans + scanl, -- :: (a -> b -> a) -> a -> Seq b -> Seq a + scanl1, -- :: (a -> a -> a) -> Seq a -> Seq a + scanr, -- :: (a -> b -> b) -> b -> Seq a -> Seq b + scanr1, -- :: (a -> a -> a) -> Seq a -> Seq a + -- * Sublists + tails, -- :: Seq a -> Seq (Seq a) + inits, -- :: Seq a -> Seq (Seq a) + -- ** Sequential searches + takeWhileL, -- :: (a -> Bool) -> Seq a -> Seq a + takeWhileR, -- :: (a -> Bool) -> Seq a -> Seq a + dropWhileL, -- :: (a -> Bool) -> Seq a -> Seq a + dropWhileR, -- :: (a -> Bool) -> Seq a -> Seq a + spanl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + spanr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + breakl, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + breakr, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + partition, -- :: (a -> Bool) -> Seq a -> (Seq a, Seq a) + filter, -- :: (a -> Bool) -> Seq a -> Seq a + -- * Sorting + sort, -- :: Ord a => Seq a -> Seq a + sortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a + unstableSort, -- :: Ord a => Seq a -> Seq a + unstableSortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a + -- * Indexing + index, -- :: Seq a -> Int -> a + adjust, -- :: (a -> a) -> Int -> Seq a -> Seq a + update, -- :: Int -> a -> Seq a -> Seq a + take, -- :: Int -> Seq a -> Seq a + drop, -- :: Int -> Seq a -> Seq a + splitAt, -- :: Int -> Seq a -> (Seq a, Seq a) + -- ** Indexing with predicates + -- | These functions perform sequential searches from the left + -- or right ends of the sequence, returning indices of matching + -- elements. + elemIndexL, -- :: Eq a => a -> Seq a -> Maybe Int + elemIndicesL, -- :: Eq a => a -> Seq a -> [Int] + elemIndexR, -- :: Eq a => a -> Seq a -> Maybe Int + elemIndicesR, -- :: Eq a => a -> Seq a -> [Int] + findIndexL, -- :: (a -> Bool) -> Seq a -> Maybe Int + findIndicesL, -- :: (a -> Bool) -> Seq a -> [Int] + findIndexR, -- :: (a -> Bool) -> Seq a -> Maybe Int + findIndicesR, -- :: (a -> Bool) -> Seq a -> [Int] + -- * Folds + -- | General folds are available via the 'Foldable' instance of 'Seq'. + foldlWithIndex, -- :: (b -> Int -> a -> b) -> b -> Seq a -> b + foldrWithIndex, -- :: (Int -> a -> b -> b) -> b -> Seq a -> b + -- * Transformations + mapWithIndex, -- :: (Int -> a -> b) -> Seq a -> Seq b + reverse, -- :: Seq a -> Seq a + -- ** Zips + zip, -- :: Seq a -> Seq b -> Seq (a, b) + zipWith, -- :: (a -> b -> c) -> Seq a -> Seq b -> Seq c + zip3, -- :: Seq a -> Seq b -> Seq c -> Seq (a, b, c) + zipWith3, -- :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d + zip4, -- :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a, b, c, d) + zipWith4, -- :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e +#if TESTING + Sized(..), + deep, + node2, + node3, +#endif + ) where + +import Prelude hiding ( + Functor(..), + null, length, take, drop, splitAt, foldl, foldl1, foldr, foldr1, + scanl, scanl1, scanr, scanr1, replicate, zip, zipWith, zip3, zipWith3, + takeWhile, dropWhile, iterate, reverse, filter, mapM, sum, all) +import qualified Data.List +import Control.Applicative (Applicative(..), (<$>), Alternative, + WrappedMonad(..), liftA, liftA2, liftA3) +import qualified Control.Applicative as Applicative (Alternative(..)) +import Control.DeepSeq (NFData(rnf)) +import Control.Monad (MonadPlus(..), ap) +import Data.Monoid (Monoid(..)) +import Data.Functor (Functor(..)) +import Data.Foldable +import Data.Traversable +import Data.Typeable + +#ifdef __GLASGOW_HASKELL__ +import GHC.Exts (build) +import Text.Read (Lexeme(Ident), lexP, parens, prec, + readPrec, readListPrec, readListPrecDefault) +import Data.Data +#endif + +infixr 5 `consTree` +infixl 5 `snocTree` + +infixr 5 >< +infixr 5 <|, :< +infixl 5 |>, :> + +class Sized a where + size :: a -> Int + +-- | General-purpose finite sequences. +newtype Seq a = Seq (FingerTree (Elem a)) + +instance Functor Seq where + fmap f (Seq xs) = Seq (fmap (fmap f) xs) +#ifdef __GLASGOW_HASKELL__ + x <$ s = replicate (length s) x +#endif + +instance Foldable Seq where + foldr f z (Seq xs) = foldr (flip (foldr f)) z xs + foldl f z (Seq xs) = foldl (foldl f) z xs + + foldr1 f (Seq xs) = getElem (foldr1 f' xs) + where f' (Elem x) (Elem y) = Elem (f x y) + + foldl1 f (Seq xs) = getElem (foldl1 f' xs) + where f' (Elem x) (Elem y) = Elem (f x y) + +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 + where add ys x = ys >< f x + +instance Applicative Seq where + pure = singleton + fs <*> xs = foldl' add empty fs + where add ys f = ys >< fmap f xs + +instance MonadPlus Seq where + mzero = empty + mplus = (><) + +instance Alternative Seq where + empty = empty + (<|>) = (><) + +instance Eq a => Eq (Seq a) where + xs == ys = length xs == length ys && toList xs == toList ys + +instance Ord a => Ord (Seq a) where + compare xs ys = compare (toList xs) (toList ys) + +#if TESTING +instance Show a => Show (Seq a) where + showsPrec p (Seq x) = showsPrec p x +#else +instance Show a => Show (Seq a) where + showsPrec p xs = showParen (p > 10) $ + showString "fromList " . shows (toList xs) +#endif + +instance Read a => Read (Seq a) where +#ifdef __GLASGOW_HASKELL__ + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- readPrec + return (fromList xs) + + readListPrec = readListPrecDefault +#else + readsPrec p = readParen (p > 10) $ \ r -> do + ("fromList",s) <- lex r + (xs,t) <- reads s + return (fromList xs,t) +#endif + +instance Monoid (Seq a) where + mempty = empty + mappend = (><) + +#include "Typeable.h" +INSTANCE_TYPEABLE1(Seq,seqTc,"Seq") + +#if __GLASGOW_HASKELL__ +instance Data a => Data (Seq a) where + gfoldl f z s = case viewl s of + EmptyL -> z empty + x :< xs -> z (<|) `f` x `f` xs + + gunfold k z c = case constrIndex c of + 1 -> z empty + 2 -> k (k (z (<|))) + _ -> error "gunfold" + + toConstr xs + | null xs = emptyConstr + | otherwise = consConstr + + dataTypeOf _ = seqDataType + + dataCast1 f = gcast1 f + +emptyConstr, consConstr :: Constr +emptyConstr = mkConstr seqDataType "empty" [] Prefix +consConstr = mkConstr seqDataType "<|" [] Infix + +seqDataType :: DataType +seqDataType = mkDataType "Data.Sequence.Seq" [emptyConstr, consConstr] +#endif + +-- Finger trees + +data FingerTree a + = Empty + | Single a + | Deep {-# UNPACK #-} !Int !(Digit a) (FingerTree (Node a)) !(Digit a) +#if TESTING + deriving Show +#endif + +instance Sized a => Sized (FingerTree a) where + {-# SPECIALIZE instance Sized (FingerTree (Elem a)) #-} + {-# SPECIALIZE instance Sized (FingerTree (Node a)) #-} + size Empty = 0 + size (Single x) = size x + size (Deep v _ _ _) = v + +instance Foldable FingerTree where + foldr _ z Empty = z + foldr f z (Single x) = x `f` z + foldr f z (Deep _ pr m sf) = + foldr f (foldr (flip (foldr f)) (foldr f z sf) m) pr + + foldl _ z Empty = z + foldl f z (Single x) = z `f` x + foldl f z (Deep _ pr m sf) = + foldl f (foldl (foldl f) (foldl f z pr) m) sf + + foldr1 _ Empty = error "foldr1: empty sequence" + foldr1 _ (Single x) = x + foldr1 f (Deep _ pr m sf) = + foldr f (foldr (flip (foldr f)) (foldr1 f sf) m) pr + + foldl1 _ Empty = error "foldl1: empty sequence" + foldl1 _ (Single x) = x + foldl1 f (Deep _ pr m sf) = + foldl f (foldl (foldl f) (foldl1 f pr) m) sf + +instance Functor FingerTree where + fmap _ Empty = Empty + fmap f (Single x) = Single (f x) + fmap f (Deep v pr m sf) = + Deep v (fmap f pr) (fmap (fmap f) m) (fmap f sf) + +instance Traversable FingerTree where + traverse _ Empty = pure Empty + traverse f (Single x) = Single <$> f x + traverse f (Deep v pr m sf) = + Deep v <$> traverse f pr <*> traverse (traverse f) m <*> + traverse f sf + +instance NFData a => NFData (FingerTree a) where + rnf (Empty) = () + rnf (Single x) = rnf x + rnf (Deep _ pr m sf) = rnf pr `seq` rnf m `seq` rnf sf + +{-# INLINE deep #-} +deep :: Sized a => Digit a -> FingerTree (Node a) -> Digit a -> FingerTree a +deep pr m sf = Deep (size pr + size m + size sf) pr m sf + +{-# INLINE pullL #-} +pullL :: Sized a => Int -> FingerTree (Node a) -> Digit a -> FingerTree a +pullL s m sf = case viewLTree m of + Nothing2 -> digitToTree' s sf + Just2 pr m' -> Deep s (nodeToDigit pr) m' sf + +{-# INLINE pullR #-} +pullR :: Sized a => Int -> Digit a -> FingerTree (Node a) -> FingerTree a +pullR s pr m = case viewRTree m of + Nothing2 -> digitToTree' s pr + Just2 m' sf -> Deep s pr m' (nodeToDigit sf) + +{-# SPECIALIZE deepL :: Maybe (Digit (Elem a)) -> FingerTree (Node (Elem a)) -> Digit (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE deepL :: Maybe (Digit (Node a)) -> FingerTree (Node (Node a)) -> Digit (Node a) -> FingerTree (Node a) #-} +deepL :: Sized a => Maybe (Digit a) -> FingerTree (Node a) -> Digit a -> FingerTree a +deepL Nothing m sf = pullL (size m + size sf) m sf +deepL (Just pr) m sf = deep pr m sf + +{-# SPECIALIZE deepR :: Digit (Elem a) -> FingerTree (Node (Elem a)) -> Maybe (Digit (Elem a)) -> FingerTree (Elem a) #-} +{-# SPECIALIZE deepR :: Digit (Node a) -> FingerTree (Node (Node a)) -> Maybe (Digit (Node a)) -> FingerTree (Node a) #-} +deepR :: Sized a => Digit a -> FingerTree (Node a) -> Maybe (Digit a) -> FingerTree a +deepR pr m Nothing = pullR (size m + size pr) pr m +deepR pr m (Just sf) = deep pr m sf + +-- Digits + +data Digit a + = One a + | Two a a + | Three a a a + | Four a a a a +#if TESTING + deriving Show +#endif + +instance Foldable Digit where + foldr f z (One a) = a `f` z + foldr f z (Two a b) = a `f` (b `f` z) + foldr f z (Three a b c) = a `f` (b `f` (c `f` z)) + foldr f z (Four a b c d) = a `f` (b `f` (c `f` (d `f` z))) + + foldl f z (One a) = z `f` a + foldl f z (Two a b) = (z `f` a) `f` b + foldl f z (Three a b c) = ((z `f` a) `f` b) `f` c + foldl f z (Four a b c d) = (((z `f` a) `f` b) `f` c) `f` d + + foldr1 _ (One a) = a + foldr1 f (Two a b) = a `f` b + foldr1 f (Three a b c) = a `f` (b `f` c) + foldr1 f (Four a b c d) = a `f` (b `f` (c `f` d)) + + foldl1 _ (One a) = a + foldl1 f (Two a b) = a `f` b + foldl1 f (Three a b c) = (a `f` b) `f` c + foldl1 f (Four a b c d) = ((a `f` b) `f` c) `f` d + +instance Functor Digit where + {-# INLINE fmap #-} + fmap f (One a) = One (f a) + fmap f (Two a b) = Two (f a) (f b) + fmap f (Three a b c) = Three (f a) (f b) (f c) + fmap f (Four a b c d) = Four (f a) (f b) (f c) (f d) + +instance Traversable Digit where + {-# INLINE traverse #-} + traverse f (One a) = One <$> f a + traverse f (Two a b) = Two <$> f a <*> f b + traverse f (Three a b c) = Three <$> f a <*> f b <*> f c + traverse f (Four a b c d) = Four <$> f a <*> f b <*> f c <*> f d + +instance NFData a => NFData (Digit a) where + rnf (One a) = rnf a + rnf (Two a b) = rnf a `seq` rnf b + rnf (Three a b c) = rnf a `seq` rnf b `seq` rnf c + rnf (Four a b c d) = rnf a `seq` rnf b `seq` rnf c `seq` rnf d + +instance Sized a => Sized (Digit a) where + {-# INLINE size #-} + size = foldl1 (+) . fmap size + +{-# SPECIALIZE digitToTree :: Digit (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE digitToTree :: Digit (Node a) -> FingerTree (Node a) #-} +digitToTree :: Sized a => Digit a -> FingerTree a +digitToTree (One a) = Single a +digitToTree (Two a b) = deep (One a) Empty (One b) +digitToTree (Three a b c) = deep (Two a b) Empty (One c) +digitToTree (Four a b c d) = deep (Two a b) Empty (Two c d) + +-- | Given the size of a digit and the digit itself, efficiently converts +-- it to a FingerTree. +digitToTree' :: Int -> Digit a -> FingerTree a +digitToTree' n (Four a b c d) = Deep n (Two a b) Empty (Two c d) +digitToTree' n (Three a b c) = Deep n (Two a b) Empty (One c) +digitToTree' n (Two a b) = Deep n (One a) Empty (One b) +digitToTree' n (One a) = n `seq` Single a + +-- Nodes + +data Node a + = Node2 {-# UNPACK #-} !Int a a + | Node3 {-# UNPACK #-} !Int a a a +#if TESTING + deriving Show +#endif + +instance Foldable Node where + foldr f z (Node2 _ a b) = a `f` (b `f` z) + foldr f z (Node3 _ a b c) = a `f` (b `f` (c `f` z)) + + foldl f z (Node2 _ a b) = (z `f` a) `f` b + foldl f z (Node3 _ a b c) = ((z `f` a) `f` b) `f` c + +instance Functor Node where + {-# INLINE fmap #-} + fmap f (Node2 v a b) = Node2 v (f a) (f b) + fmap f (Node3 v a b c) = Node3 v (f a) (f b) (f c) + +instance Traversable Node where + {-# INLINE traverse #-} + traverse f (Node2 v a b) = Node2 v <$> f a <*> f b + traverse f (Node3 v a b c) = Node3 v <$> f a <*> f b <*> f c + +instance NFData a => NFData (Node a) where + rnf (Node2 _ a b) = rnf a `seq` rnf b + rnf (Node3 _ a b c) = rnf a `seq` rnf b `seq` rnf c + +instance Sized (Node a) where + size (Node2 v _ _) = v + size (Node3 v _ _ _) = v + +{-# INLINE node2 #-} +node2 :: Sized a => a -> a -> Node a +node2 a b = Node2 (size a + size b) a b + +{-# INLINE node3 #-} +node3 :: Sized a => a -> a -> a -> Node a +node3 a b c = Node3 (size a + size b + size c) a b c + +nodeToDigit :: Node a -> Digit a +nodeToDigit (Node2 _ a b) = Two a b +nodeToDigit (Node3 _ a b c) = Three a b c + +-- Elements + +newtype Elem a = Elem { getElem :: a } +#if TESTING + deriving Show +#endif + +instance Sized (Elem a) where + size _ = 1 + +instance Functor Elem where + fmap f (Elem x) = Elem (f x) + +instance Foldable Elem where + foldr f z (Elem x) = f x z + foldl f z (Elem x) = f z x + +instance Traversable Elem where + traverse f (Elem x) = Elem <$> f x + +instance NFData a => NFData (Elem a) where + rnf (Elem x) = rnf x + +------------------------------------------------------- +-- Applicative construction +------------------------------------------------------- + +newtype Id a = Id {runId :: a} + +instance Functor Id where + fmap f (Id x) = Id (f x) + +instance Monad Id where + return = Id + m >>= k = k (runId m) + +instance Applicative Id where + pure = return + (<*>) = ap + +-- | This is essentially a clone of Control.Monad.State.Strict. +newtype State s a = State {runState :: s -> (s, a)} + +instance Functor (State s) where + fmap = liftA + +instance Monad (State s) where + {-# INLINE return #-} + {-# INLINE (>>=) #-} + return x = State $ \ s -> (s, x) + m >>= k = State $ \ s -> case runState m s of + (s', x) -> runState (k x) s' + +instance Applicative (State s) where + pure = return + (<*>) = ap + +execState :: State s a -> s -> a +execState m x = snd (runState m x) + +-- | A helper method: a strict version of mapAccumL. +mapAccumL' :: Traversable t => (a -> b -> (a, c)) -> a -> t b -> (a, t c) +mapAccumL' f s t = runState (traverse (State . flip f) t) s + +-- | 'applicativeTree' takes an Applicative-wrapped construction of a +-- piece of a FingerTree, assumed to always have the same size (which +-- is put in the second argument), and replicates it as many times as +-- specified. This is a generalization of 'replicateA', which itself +-- is a generalization of many Data.Sequence methods. +{-# SPECIALIZE applicativeTree :: Int -> Int -> State s a -> State s (FingerTree a) #-} +{-# SPECIALIZE applicativeTree :: Int -> Int -> Id a -> Id (FingerTree a) #-} +-- Special note: the Id specialization automatically does node sharing, +-- reducing memory usage of the resulting tree to /O(log n)/. +applicativeTree :: Applicative f => Int -> Int -> f a -> f (FingerTree a) +applicativeTree n mSize m = mSize `seq` case n of + 0 -> pure Empty + 1 -> liftA Single m + 2 -> deepA one emptyTree one + 3 -> deepA two emptyTree one + 4 -> deepA two emptyTree two + 5 -> deepA three emptyTree two + 6 -> deepA three emptyTree three + 7 -> deepA four emptyTree three + 8 -> deepA four emptyTree four + _ -> let (q, r) = n `quotRem` 3 in q `seq` case r of + 0 -> deepA three (applicativeTree (q - 2) mSize' n3) three + 1 -> deepA four (applicativeTree (q - 2) mSize' n3) three + _ -> deepA four (applicativeTree (q - 2) mSize' n3) four + where + one = liftA One m + two = liftA2 Two m m + three = liftA3 Three m m m + four = liftA3 Four m m m <*> m + deepA = liftA3 (Deep (n * mSize)) + mSize' = 3 * mSize + n3 = liftA3 (Node3 mSize') m m m + emptyTree = pure Empty + +------------------------------------------------------------------------ +-- Construction +------------------------------------------------------------------------ + +-- | /O(1)/. The empty sequence. +empty :: Seq a +empty = Seq Empty + +-- | /O(1)/. A singleton sequence. +singleton :: a -> Seq a +singleton x = Seq (Single (Elem x)) + +-- | /O(log n)/. @replicate n x@ is a sequence consisting of @n@ copies of @x@. +replicate :: Int -> a -> Seq a +replicate n x + | n >= 0 = runId (replicateA n (Id x)) + | otherwise = error "replicate takes a nonnegative integer argument" + +-- | 'replicateA' is an 'Applicative' version of 'replicate', and makes +-- /O(log n)/ calls to '<*>' and 'pure'. +-- +-- > replicateA n x = sequenceA (replicate n x) +replicateA :: Applicative f => Int -> f a -> f (Seq a) +replicateA n x + | n >= 0 = Seq <$> applicativeTree n 1 (Elem <$> x) + | otherwise = error "replicateA takes a nonnegative integer argument" + +-- | 'replicateM' is a sequence counterpart of 'Control.Monad.replicateM'. +-- +-- > replicateM n x = sequence (replicate n x) +replicateM :: Monad m => Int -> m a -> m (Seq a) +replicateM n x + | n >= 0 = unwrapMonad (replicateA n (WrapMonad x)) + | otherwise = error "replicateM takes a nonnegative integer argument" + +-- | /O(1)/. Add an element to the left end of a sequence. +-- Mnemonic: a triangle with the single element at the pointy end. +(<|) :: a -> Seq a -> Seq a +x <| Seq xs = Seq (Elem x `consTree` xs) + +{-# SPECIALIZE consTree :: Elem a -> FingerTree (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE consTree :: Node a -> FingerTree (Node a) -> FingerTree (Node a) #-} +consTree :: Sized a => a -> FingerTree a -> FingerTree a +consTree a Empty = Single a +consTree a (Single b) = deep (One a) Empty (One b) +consTree a (Deep s (Four b c d e) m sf) = m `seq` + Deep (size a + s) (Two a b) (node3 c d e `consTree` m) sf +consTree a (Deep s (Three b c d) m sf) = + Deep (size a + s) (Four a b c d) m sf +consTree a (Deep s (Two b c) m sf) = + Deep (size a + s) (Three a b c) m sf +consTree a (Deep s (One b) m sf) = + Deep (size a + s) (Two a b) m sf + +-- | /O(1)/. Add an element to the right end of a sequence. +-- Mnemonic: a triangle with the single element at the pointy end. +(|>) :: Seq a -> a -> Seq a +Seq xs |> x = Seq (xs `snocTree` Elem x) + +{-# SPECIALIZE snocTree :: FingerTree (Elem a) -> Elem a -> FingerTree (Elem a) #-} +{-# SPECIALIZE snocTree :: FingerTree (Node a) -> Node a -> FingerTree (Node a) #-} +snocTree :: Sized a => FingerTree a -> a -> FingerTree a +snocTree Empty a = Single a +snocTree (Single a) b = deep (One a) Empty (One b) +snocTree (Deep s pr m (Four a b c d)) e = m `seq` + Deep (s + size e) pr (m `snocTree` node3 a b c) (Two d e) +snocTree (Deep s pr m (Three a b c)) d = + Deep (s + size d) pr m (Four a b c d) +snocTree (Deep s pr m (Two a b)) c = + Deep (s + size c) pr m (Three a b c) +snocTree (Deep s pr m (One a)) b = + Deep (s + size b) pr m (Two a b) + +-- | /O(log(min(n1,n2)))/. Concatenate two sequences. +(><) :: Seq a -> Seq a -> Seq a +Seq xs >< Seq ys = Seq (appendTree0 xs ys) + +-- The appendTree/addDigits gunk below is machine generated + +appendTree0 :: FingerTree (Elem a) -> FingerTree (Elem a) -> FingerTree (Elem a) +appendTree0 Empty xs = + xs +appendTree0 xs Empty = + xs +appendTree0 (Single x) xs = + x `consTree` xs +appendTree0 xs (Single x) = + xs `snocTree` x +appendTree0 (Deep s1 pr1 m1 sf1) (Deep s2 pr2 m2 sf2) = + Deep (s1 + s2) pr1 (addDigits0 m1 sf1 pr2 m2) sf2 + +addDigits0 :: FingerTree (Node (Elem a)) -> Digit (Elem a) -> Digit (Elem a) -> FingerTree (Node (Elem a)) -> FingerTree (Node (Elem a)) +addDigits0 m1 (One a) (One b) m2 = + appendTree1 m1 (node2 a b) m2 +addDigits0 m1 (One a) (Two b c) m2 = + appendTree1 m1 (node3 a b c) m2 +addDigits0 m1 (One a) (Three b c d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits0 m1 (One a) (Four b c d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Two a b) (One c) m2 = + appendTree1 m1 (node3 a b c) m2 +addDigits0 m1 (Two a b) (Two c d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits0 m1 (Two a b) (Three c d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Two a b) (Four c d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits0 m1 (Three a b c) (One d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits0 m1 (Three a b c) (Two d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Three a b c) (Three d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits0 m1 (Three a b c) (Four d e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits0 m1 (Four a b c d) (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits0 m1 (Four a b c d) (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits0 m1 (Four a b c d) (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits0 m1 (Four a b c d) (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 + +appendTree1 :: FingerTree (Node a) -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree1 Empty a xs = + a `consTree` xs +appendTree1 xs a Empty = + xs `snocTree` a +appendTree1 (Single x) a xs = + x `consTree` a `consTree` xs +appendTree1 xs a (Single x) = + xs `snocTree` a `snocTree` x +appendTree1 (Deep s1 pr1 m1 sf1) a (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + s2) pr1 (addDigits1 m1 sf1 a pr2 m2) sf2 + +addDigits1 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits1 m1 (One a) b (One c) m2 = + appendTree1 m1 (node3 a b c) m2 +addDigits1 m1 (One a) b (Two c d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits1 m1 (One a) b (Three c d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits1 m1 (One a) b (Four c d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Two a b) c (One d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits1 m1 (Two a b) c (Two d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits1 m1 (Two a b) c (Three d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Two a b) c (Four d e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits1 m1 (Three a b c) d (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits1 m1 (Three a b c) d (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Three a b c) d (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits1 m1 (Three a b c) d (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits1 m1 (Four a b c d) e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits1 m1 (Four a b c d) e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits1 m1 (Four a b c d) e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits1 m1 (Four a b c d) e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 + +appendTree2 :: FingerTree (Node a) -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree2 Empty a b xs = + a `consTree` b `consTree` xs +appendTree2 xs a b Empty = + xs `snocTree` a `snocTree` b +appendTree2 (Single x) a b xs = + x `consTree` a `consTree` b `consTree` xs +appendTree2 xs a b (Single x) = + xs `snocTree` a `snocTree` b `snocTree` x +appendTree2 (Deep s1 pr1 m1 sf1) a b (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + size b + s2) pr1 (addDigits2 m1 sf1 a b pr2 m2) sf2 + +addDigits2 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits2 m1 (One a) b c (One d) m2 = + appendTree2 m1 (node2 a b) (node2 c d) m2 +addDigits2 m1 (One a) b c (Two d e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits2 m1 (One a) b c (Three d e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits2 m1 (One a) b c (Four d e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Two a b) c d (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits2 m1 (Two a b) c d (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits2 m1 (Two a b) c d (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Two a b) c d (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits2 m1 (Three a b c) d e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits2 m1 (Three a b c) d e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Three a b c) d e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits2 m1 (Three a b c) d e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits2 m1 (Four a b c d) e f (One g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits2 m1 (Four a b c d) e f (Two g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits2 m1 (Four a b c d) e f (Three g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits2 m1 (Four a b c d) e f (Four g h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 + +appendTree3 :: FingerTree (Node a) -> Node a -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree3 Empty a b c xs = + a `consTree` b `consTree` c `consTree` xs +appendTree3 xs a b c Empty = + xs `snocTree` a `snocTree` b `snocTree` c +appendTree3 (Single x) a b c xs = + x `consTree` a `consTree` b `consTree` c `consTree` xs +appendTree3 xs a b c (Single x) = + xs `snocTree` a `snocTree` b `snocTree` c `snocTree` x +appendTree3 (Deep s1 pr1 m1 sf1) a b c (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + size b + size c + s2) pr1 (addDigits3 m1 sf1 a b c pr2 m2) sf2 + +addDigits3 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits3 m1 (One a) b c d (One e) m2 = + appendTree2 m1 (node3 a b c) (node2 d e) m2 +addDigits3 m1 (One a) b c d (Two e f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits3 m1 (One a) b c d (Three e f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits3 m1 (One a) b c d (Four e f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Two a b) c d e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits3 m1 (Two a b) c d e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits3 m1 (Two a b) c d e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Two a b) c d e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits3 m1 (Three a b c) d e f (One g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits3 m1 (Three a b c) d e f (Two g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Three a b c) d e f (Three g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits3 m1 (Three a b c) d e f (Four g h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits3 m1 (Four a b c d) e f g (One h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits3 m1 (Four a b c d) e f g (Two h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits3 m1 (Four a b c d) e f g (Three h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits3 m1 (Four a b c d) e f g (Four h i j k) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 + +appendTree4 :: FingerTree (Node a) -> Node a -> Node a -> Node a -> Node a -> FingerTree (Node a) -> FingerTree (Node a) +appendTree4 Empty a b c d xs = + a `consTree` b `consTree` c `consTree` d `consTree` xs +appendTree4 xs a b c d Empty = + xs `snocTree` a `snocTree` b `snocTree` c `snocTree` d +appendTree4 (Single x) a b c d xs = + x `consTree` a `consTree` b `consTree` c `consTree` d `consTree` xs +appendTree4 xs a b c d (Single x) = + xs `snocTree` a `snocTree` b `snocTree` c `snocTree` d `snocTree` x +appendTree4 (Deep s1 pr1 m1 sf1) a b c d (Deep s2 pr2 m2 sf2) = + Deep (s1 + size a + size b + size c + size d + s2) pr1 (addDigits4 m1 sf1 a b c d pr2 m2) sf2 + +addDigits4 :: FingerTree (Node (Node a)) -> Digit (Node a) -> Node a -> Node a -> Node a -> Node a -> Digit (Node a) -> FingerTree (Node (Node a)) -> FingerTree (Node (Node a)) +addDigits4 m1 (One a) b c d e (One f) m2 = + appendTree2 m1 (node3 a b c) (node3 d e f) m2 +addDigits4 m1 (One a) b c d e (Two f g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits4 m1 (One a) b c d e (Three f g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits4 m1 (One a) b c d e (Four f g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Two a b) c d e f (One g) m2 = + appendTree3 m1 (node3 a b c) (node2 d e) (node2 f g) m2 +addDigits4 m1 (Two a b) c d e f (Two g h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits4 m1 (Two a b) c d e f (Three g h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Two a b) c d e f (Four g h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits4 m1 (Three a b c) d e f g (One h) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node2 g h) m2 +addDigits4 m1 (Three a b c) d e f g (Two h i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Three a b c) d e f g (Three h i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits4 m1 (Three a b c) d e f g (Four h i j k) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 +addDigits4 m1 (Four a b c d) e f g h (One i) m2 = + appendTree3 m1 (node3 a b c) (node3 d e f) (node3 g h i) m2 +addDigits4 m1 (Four a b c d) e f g h (Two i j) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node2 g h) (node2 i j) m2 +addDigits4 m1 (Four a b c d) e f g h (Three i j k) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node2 j k) m2 +addDigits4 m1 (Four a b c d) e f g h (Four i j k l) m2 = + appendTree4 m1 (node3 a b c) (node3 d e f) (node3 g h i) (node3 j k l) m2 + +-- | Builds a sequence from a seed value. Takes time linear in the +-- number of generated elements. /WARNING:/ If the number of generated +-- elements is infinite, this method will not terminate. +unfoldr :: (b -> Maybe (a, b)) -> b -> Seq a +unfoldr f = unfoldr' empty + -- uses tail recursion rather than, for instance, the List implementation. + where unfoldr' as b = maybe as (\ (a, b') -> unfoldr' (as |> a) b') (f b) + +-- | @'unfoldl' f x@ is equivalent to @'reverse' ('unfoldr' ('fmap' swap . f) x)@. +unfoldl :: (b -> Maybe (b, a)) -> b -> Seq a +unfoldl f = unfoldl' empty + where unfoldl' as b = maybe as (\ (b', a) -> unfoldl' (a <| as) b') (f b) + +-- | /O(n)/. Constructs a sequence by repeated application of a function +-- to a seed value. +-- +-- > iterateN n f x = fromList (Prelude.take n (Prelude.iterate f x)) +iterateN :: Int -> (a -> a) -> a -> Seq a +iterateN n f x + | n >= 0 = replicateA n (State (\ y -> (f y, y))) `execState` x + | otherwise = error "iterateN takes a nonnegative integer argument" + +------------------------------------------------------------------------ +-- Deconstruction +------------------------------------------------------------------------ + +-- | /O(1)/. Is this the empty sequence? +null :: Seq a -> Bool +null (Seq Empty) = True +null _ = False + +-- | /O(1)/. The number of elements in the sequence. +length :: Seq a -> Int +length (Seq xs) = size xs + +-- Views + +data Maybe2 a b = Nothing2 | Just2 a b + +-- | View of the left end of a sequence. +data ViewL a + = EmptyL -- ^ empty sequence + | a :< Seq a -- ^ leftmost element and the rest of the sequence +#if __GLASGOW_HASKELL__ + deriving (Eq, Ord, Show, Read, Data) +#else + deriving (Eq, Ord, Show, Read) +#endif + +INSTANCE_TYPEABLE1(ViewL,viewLTc,"ViewL") + +instance Functor ViewL where + {-# INLINE fmap #-} + fmap _ EmptyL = EmptyL + fmap f (x :< xs) = f x :< fmap f xs + +instance Foldable ViewL where + foldr _ z EmptyL = z + foldr f z (x :< xs) = f x (foldr f z xs) + + foldl _ z EmptyL = z + foldl f z (x :< xs) = foldl f (f z x) xs + + foldl1 _ EmptyL = error "foldl1: empty view" + foldl1 f (x :< xs) = foldl f x xs + +instance Traversable ViewL where + traverse _ EmptyL = pure EmptyL + traverse f (x :< xs) = (:<) <$> f x <*> traverse f xs + +-- | /O(1)/. Analyse the left end of a sequence. +viewl :: Seq a -> ViewL a +viewl (Seq xs) = case viewLTree xs of + Nothing2 -> EmptyL + Just2 (Elem x) xs' -> x :< Seq xs' + +{-# SPECIALIZE viewLTree :: FingerTree (Elem a) -> Maybe2 (Elem a) (FingerTree (Elem a)) #-} +{-# SPECIALIZE viewLTree :: FingerTree (Node a) -> Maybe2 (Node a) (FingerTree (Node a)) #-} +viewLTree :: Sized a => FingerTree a -> Maybe2 a (FingerTree a) +viewLTree Empty = Nothing2 +viewLTree (Single a) = Just2 a Empty +viewLTree (Deep s (One a) m sf) = Just2 a (pullL (s - size a) m sf) +viewLTree (Deep s (Two a b) m sf) = + Just2 a (Deep (s - size a) (One b) m sf) +viewLTree (Deep s (Three a b c) m sf) = + Just2 a (Deep (s - size a) (Two b c) m sf) +viewLTree (Deep s (Four a b c d) m sf) = + Just2 a (Deep (s - size a) (Three b c d) m sf) + +-- | View of the right end of a sequence. +data ViewR a + = EmptyR -- ^ empty sequence + | Seq a :> a -- ^ the sequence minus the rightmost element, + -- and the rightmost element +#if __GLASGOW_HASKELL__ + deriving (Eq, Ord, Show, Read, Data) +#else + deriving (Eq, Ord, Show, Read) +#endif + +INSTANCE_TYPEABLE1(ViewR,viewRTc,"ViewR") + +instance Functor ViewR where + {-# INLINE fmap #-} + fmap _ EmptyR = EmptyR + fmap f (xs :> x) = fmap f xs :> f x + +instance Foldable ViewR where + foldr _ z EmptyR = z + foldr f z (xs :> x) = foldr f (f x z) xs + + foldl _ z EmptyR = z + foldl f z (xs :> x) = foldl f z xs `f` x + + foldr1 _ EmptyR = error "foldr1: empty view" + foldr1 f (xs :> x) = foldr f x xs + +instance Traversable ViewR where + traverse _ EmptyR = pure EmptyR + traverse f (xs :> x) = (:>) <$> traverse f xs <*> f x + +-- | /O(1)/. Analyse the right end of a sequence. +viewr :: Seq a -> ViewR a +viewr (Seq xs) = case viewRTree xs of + Nothing2 -> EmptyR + Just2 xs' (Elem x) -> Seq xs' :> x + +{-# SPECIALIZE viewRTree :: FingerTree (Elem a) -> Maybe2 (FingerTree (Elem a)) (Elem a) #-} +{-# SPECIALIZE viewRTree :: FingerTree (Node a) -> Maybe2 (FingerTree (Node a)) (Node a) #-} +viewRTree :: Sized a => FingerTree a -> Maybe2 (FingerTree a) a +viewRTree Empty = Nothing2 +viewRTree (Single z) = Just2 Empty z +viewRTree (Deep s pr m (One z)) = Just2 (pullR (s - size z) pr m) z +viewRTree (Deep s pr m (Two y z)) = + Just2 (Deep (s - size z) pr m (One y)) z +viewRTree (Deep s pr m (Three x y z)) = + Just2 (Deep (s - size z) pr m (Two x y)) z +viewRTree (Deep s pr m (Four w x y z)) = + Just2 (Deep (s - size z) pr m (Three w x y)) z + +------------------------------------------------------------------------ +-- Scans +-- +-- These are not particularly complex applications of the Traversable +-- functor, though making the correspondence with Data.List exact +-- requires the use of (<|) and (|>). +-- +-- Note that save for the single (<|) or (|>), we maintain the original +-- structure of the Seq, not having to do any restructuring of our own. +-- +-- wasserman.louis@gmail.com, 5/23/09 +------------------------------------------------------------------------ + +-- | 'scanl' is similar to 'foldl', but returns a sequence of reduced +-- values from the left: +-- +-- > scanl f z (fromList [x1, x2, ...]) = fromList [z, z `f` x1, (z `f` x1) `f` x2, ...] +scanl :: (a -> b -> a) -> a -> Seq b -> Seq a +scanl f z0 xs = z0 <| snd (mapAccumL (\ x z -> let x' = f x z in (x', x')) z0 xs) + +-- | 'scanl1' is a variant of 'scanl' that has no starting value argument: +-- +-- > scanl1 f (fromList [x1, x2, ...]) = fromList [x1, x1 `f` x2, ...] +scanl1 :: (a -> a -> a) -> Seq a -> Seq a +scanl1 f xs = case viewl xs of + EmptyL -> error "scanl1 takes a nonempty sequence as an argument" + x :< xs' -> scanl f x xs' + +-- | 'scanr' is the right-to-left dual of 'scanl'. +scanr :: (a -> b -> b) -> b -> Seq a -> Seq b +scanr f z0 xs = snd (mapAccumR (\ z x -> let z' = f x z in (z', z')) z0 xs) |> z0 + +-- | 'scanr1' is a variant of 'scanr' that has no starting value argument. +scanr1 :: (a -> a -> a) -> Seq a -> Seq a +scanr1 f xs = case viewr xs of + EmptyR -> error "scanr1 takes a nonempty sequence as an argument" + xs' :> x -> scanr f x xs' + +-- Indexing + +-- | /O(log(min(i,n-i)))/. The element at the specified position, +-- counting from 0. The argument should thus be a non-negative +-- integer less than the size of the sequence. +-- If the position is out of range, 'index' fails with an error. +index :: Seq a -> Int -> a +index (Seq xs) i + | 0 <= i && i < size xs = case lookupTree i xs of + Place _ (Elem x) -> x + | otherwise = error "index out of bounds" + +data Place a = Place {-# UNPACK #-} !Int a +#if TESTING + deriving Show +#endif + +{-# SPECIALIZE lookupTree :: Int -> FingerTree (Elem a) -> Place (Elem a) #-} +{-# SPECIALIZE lookupTree :: Int -> FingerTree (Node a) -> Place (Node a) #-} +lookupTree :: Sized a => Int -> FingerTree a -> Place a +lookupTree _ Empty = error "lookupTree of empty tree" +lookupTree i (Single x) = Place i x +lookupTree i (Deep _ pr m sf) + | i < spr = lookupDigit i pr + | i < spm = case lookupTree (i - spr) m of + Place i' xs -> lookupNode i' xs + | otherwise = lookupDigit (i - spm) sf + where + spr = size pr + spm = spr + size m + +{-# SPECIALIZE lookupNode :: Int -> Node (Elem a) -> Place (Elem a) #-} +{-# SPECIALIZE lookupNode :: Int -> Node (Node a) -> Place (Node a) #-} +lookupNode :: Sized a => Int -> Node a -> Place a +lookupNode i (Node2 _ a b) + | i < sa = Place i a + | otherwise = Place (i - sa) b + where + sa = size a +lookupNode i (Node3 _ a b c) + | i < sa = Place i a + | i < sab = Place (i - sa) b + | otherwise = Place (i - sab) c + where + sa = size a + sab = sa + size b + +{-# SPECIALIZE lookupDigit :: Int -> Digit (Elem a) -> Place (Elem a) #-} +{-# SPECIALIZE lookupDigit :: Int -> Digit (Node a) -> Place (Node a) #-} +lookupDigit :: Sized a => Int -> Digit a -> Place a +lookupDigit i (One a) = Place i a +lookupDigit i (Two a b) + | i < sa = Place i a + | otherwise = Place (i - sa) b + where + sa = size a +lookupDigit i (Three a b c) + | i < sa = Place i a + | i < sab = Place (i - sa) b + | otherwise = Place (i - sab) c + where + sa = size a + sab = sa + size b +lookupDigit i (Four a b c d) + | i < sa = Place i a + | i < sab = Place (i - sa) b + | i < sabc = Place (i - sab) c + | otherwise = Place (i - sabc) d + where + sa = size a + sab = sa + size b + sabc = sab + size c + +-- | /O(log(min(i,n-i)))/. Replace the element at the specified position. +-- If the position is out of range, the original sequence is returned. +update :: Int -> a -> Seq a -> Seq a +update i x = adjust (const x) i + +-- | /O(log(min(i,n-i)))/. Update the element at the specified position. +-- If the position is out of range, the original sequence is returned. +adjust :: (a -> a) -> Int -> Seq a -> Seq a +adjust f i (Seq xs) + | 0 <= i && i < size xs = Seq (adjustTree (const (fmap f)) i xs) + | otherwise = Seq xs + +{-# SPECIALIZE adjustTree :: (Int -> Elem a -> Elem a) -> Int -> FingerTree (Elem a) -> FingerTree (Elem a) #-} +{-# SPECIALIZE adjustTree :: (Int -> Node a -> Node a) -> Int -> FingerTree (Node a) -> FingerTree (Node a) #-} +adjustTree :: Sized a => (Int -> a -> a) -> + Int -> FingerTree a -> FingerTree a +adjustTree _ _ Empty = error "adjustTree of empty tree" +adjustTree f i (Single x) = Single (f i x) +adjustTree f i (Deep s pr m sf) + | i < spr = Deep s (adjustDigit f i pr) m sf + | i < spm = Deep s pr (adjustTree (adjustNode f) (i - spr) m) sf + | otherwise = Deep s pr m (adjustDigit f (i - spm) sf) + where + spr = size pr + spm = spr + size m + +{-# SPECIALIZE adjustNode :: (Int -> Elem a -> Elem a) -> Int -> Node (Elem a) -> Node (Elem a) #-} +{-# SPECIALIZE adjustNode :: (Int -> Node a -> Node a) -> Int -> Node (Node a) -> Node (Node a) #-} +adjustNode :: Sized a => (Int -> a -> a) -> Int -> Node a -> Node a +adjustNode f i (Node2 s a b) + | i < sa = Node2 s (f i a) b + | otherwise = Node2 s a (f (i - sa) b) + where + sa = size a +adjustNode f i (Node3 s a b c) + | i < sa = Node3 s (f i a) b c + | i < sab = Node3 s a (f (i - sa) b) c + | otherwise = Node3 s a b (f (i - sab) c) + where + sa = size a + sab = sa + size b + +{-# SPECIALIZE adjustDigit :: (Int -> Elem a -> Elem a) -> Int -> Digit (Elem a) -> Digit (Elem a) #-} +{-# SPECIALIZE adjustDigit :: (Int -> Node a -> Node a) -> Int -> Digit (Node a) -> Digit (Node a) #-} +adjustDigit :: Sized a => (Int -> a -> a) -> Int -> Digit a -> Digit a +adjustDigit f i (One a) = One (f i a) +adjustDigit f i (Two a b) + | i < sa = Two (f i a) b + | otherwise = Two a (f (i - sa) b) + where + sa = size a +adjustDigit f i (Three a b c) + | i < sa = Three (f i a) b c + | i < sab = Three a (f (i - sa) b) c + | otherwise = Three a b (f (i - sab) c) + where + sa = size a + sab = sa + size b +adjustDigit f i (Four a b c d) + | i < sa = Four (f i a) b c d + | i < sab = Four a (f (i - sa) b) c d + | i < sabc = Four a b (f (i - sab) c) d + | otherwise = Four a b c (f (i- sabc) d) + where + sa = size a + sab = sa + size b + sabc = sab + size c + +-- | A generalization of 'fmap', 'mapWithIndex' takes a mapping function +-- that also depends on the element's index, and applies it to every +-- element in the sequence. +mapWithIndex :: (Int -> a -> b) -> Seq a -> Seq b +mapWithIndex f xs = snd (mapAccumL' (\ i x -> (i + 1, f i x)) 0 xs) + +-- Splitting + +-- | /O(log(min(i,n-i)))/. The first @i@ elements of a sequence. +-- If @i@ is negative, @'take' i s@ yields the empty sequence. +-- If the sequence contains fewer than @i@ elements, the whole sequence +-- is returned. +take :: Int -> Seq a -> Seq a +take i = fst . splitAt i + +-- | /O(log(min(i,n-i)))/. Elements of a sequence after the first @i@. +-- If @i@ is negative, @'drop' i s@ yields the whole sequence. +-- If the sequence contains fewer than @i@ elements, the empty sequence +-- is returned. +drop :: Int -> Seq a -> Seq a +drop i = snd . splitAt i + +-- | /O(log(min(i,n-i)))/. Split a sequence at a given position. +-- @'splitAt' i s = ('take' i s, 'drop' i s)@. +splitAt :: Int -> Seq a -> (Seq a, Seq a) +splitAt i (Seq xs) = (Seq l, Seq r) + where (l, r) = split i xs + +split :: Int -> FingerTree (Elem a) -> + (FingerTree (Elem a), FingerTree (Elem a)) +split i Empty = i `seq` (Empty, Empty) +split i xs + | size xs > i = (l, consTree x r) + | otherwise = (xs, Empty) + where Split l x r = splitTree i xs + +data Split t a = Split t a t +#if TESTING + deriving Show +#endif + +{-# SPECIALIZE splitTree :: Int -> FingerTree (Elem a) -> Split (FingerTree (Elem a)) (Elem a) #-} +{-# SPECIALIZE splitTree :: Int -> FingerTree (Node a) -> Split (FingerTree (Node a)) (Node a) #-} +splitTree :: Sized a => Int -> FingerTree a -> Split (FingerTree a) a +splitTree _ Empty = error "splitTree of empty tree" +splitTree i (Single x) = i `seq` Split Empty x Empty +splitTree i (Deep _ pr m sf) + | i < spr = case splitDigit i pr of + Split l x r -> Split (maybe Empty digitToTree l) x (deepL r m sf) + | i < spm = case splitTree im m of + Split ml xs mr -> case splitNode (im - size ml) xs of + Split l x r -> Split (deepR pr ml l) x (deepL r mr sf) + | otherwise = case splitDigit (i - spm) sf of + Split l x r -> Split (deepR pr m l) x (maybe Empty digitToTree r) + where + spr = size pr + spm = spr + size m + im = i - spr + +{-# SPECIALIZE splitNode :: Int -> Node (Elem a) -> Split (Maybe (Digit (Elem a))) (Elem a) #-} +{-# SPECIALIZE splitNode :: Int -> Node (Node a) -> Split (Maybe (Digit (Node a))) (Node a) #-} +splitNode :: Sized a => Int -> Node a -> Split (Maybe (Digit a)) a +splitNode i (Node2 _ a b) + | i < sa = Split Nothing a (Just (One b)) + | otherwise = Split (Just (One a)) b Nothing + where + sa = size a +splitNode i (Node3 _ a b c) + | i < sa = Split Nothing a (Just (Two b c)) + | i < sab = Split (Just (One a)) b (Just (One c)) + | otherwise = Split (Just (Two a b)) c Nothing + where + sa = size a + sab = sa + size b + +{-# SPECIALIZE splitDigit :: Int -> Digit (Elem a) -> Split (Maybe (Digit (Elem a))) (Elem a) #-} +{-# SPECIALIZE splitDigit :: Int -> Digit (Node a) -> Split (Maybe (Digit (Node a))) (Node a) #-} +splitDigit :: Sized a => Int -> Digit a -> Split (Maybe (Digit a)) a +splitDigit i (One a) = i `seq` Split Nothing a Nothing +splitDigit i (Two a b) + | i < sa = Split Nothing a (Just (One b)) + | otherwise = Split (Just (One a)) b Nothing + where + sa = size a +splitDigit i (Three a b c) + | i < sa = Split Nothing a (Just (Two b c)) + | i < sab = Split (Just (One a)) b (Just (One c)) + | otherwise = Split (Just (Two a b)) c Nothing + where + sa = size a + sab = sa + size b +splitDigit i (Four a b c d) + | i < sa = Split Nothing a (Just (Three b c d)) + | i < sab = Split (Just (One a)) b (Just (Two c d)) + | i < sabc = Split (Just (Two a b)) c (Just (One d)) + | otherwise = Split (Just (Three a b c)) d Nothing + where + sa = size a + sab = sa + size b + sabc = sab + size c + +-- | /O(n)/. Returns a sequence of all suffixes of this sequence, +-- longest first. For example, +-- +-- > tails (fromList "abc") = fromList [fromList "abc", fromList "bc", fromList "c", fromList ""] +-- +-- Evaluating the /i/th suffix takes /O(log(min(i, n-i)))/, but evaluating +-- every suffix in the sequence takes /O(n)/ due to sharing. +tails :: Seq a -> Seq (Seq a) +tails (Seq xs) = Seq (tailsTree (Elem . Seq) xs) |> empty + +-- | /O(n)/. Returns a sequence of all prefixes of this sequence, +-- shortest first. For example, +-- +-- > inits (fromList "abc") = fromList [fromList "", fromList "a", fromList "ab", fromList "abc"] +-- +-- Evaluating the /i/th prefix takes /O(log(min(i, n-i)))/, but evaluating +-- every prefix in the sequence takes /O(n)/ due to sharing. +inits :: Seq a -> Seq (Seq a) +inits (Seq xs) = empty <| Seq (initsTree (Elem . Seq) xs) + +-- This implementation of tails (and, analogously, inits) has the +-- following algorithmic advantages: +-- Evaluating each tail in the sequence takes linear total time, +-- which is better than we could say for +-- @fromList [drop n xs | n <- [0..length xs]]@. +-- Evaluating any individual tail takes logarithmic time, which is +-- better than we can say for either +-- @scanr (<|) empty xs@ or @iterateN (length xs + 1) (\ xs -> let _ :< xs' = viewl xs in xs') xs@. +-- +-- Moreover, if we actually look at every tail in the sequence, the +-- following benchmarks demonstrate that this implementation is modestly +-- faster than any of the above: +-- +-- Times (ms) +-- min mean +/-sd median max +-- Seq.tails: 21.986 24.961 10.169 22.417 86.485 +-- scanr: 85.392 87.942 2.488 87.425 100.217 +-- iterateN: 29.952 31.245 1.574 30.412 37.268 +-- +-- The algorithm for tails (and, analogously, inits) is as follows: +-- +-- A Node in the FingerTree of tails is constructed by evaluating the +-- corresponding tail of the FingerTree of Nodes, considering the first +-- Node in this tail, and constructing a Node in which each tail of this +-- Node is made to be the prefix of the remaining tree. This ends up +-- working quite elegantly, as the remainder of the tail of the FingerTree +-- of Nodes becomes the middle of a new tail, the suffix of the Node is +-- the prefix, and the suffix of the original tree is retained. +-- +-- In particular, evaluating the /i/th tail involves making as +-- many partial evaluations as the Node depth of the /i/th element. +-- In addition, when we evaluate the /i/th tail, and we also evaluate +-- the /j/th tail, and /m/ Nodes are on the path to both /i/ and /j/, +-- each of those /m/ evaluations are shared between the computation of +-- the /i/th and /j/th tails. +-- +-- wasserman.louis@gmail.com, 7/16/09 + +tailsDigit :: Digit a -> Digit (Digit a) +tailsDigit (One a) = One (One a) +tailsDigit (Two a b) = Two (Two a b) (One b) +tailsDigit (Three a b c) = Three (Three a b c) (Two b c) (One c) +tailsDigit (Four a b c d) = Four (Four a b c d) (Three b c d) (Two c d) (One d) + +initsDigit :: Digit a -> Digit (Digit a) +initsDigit (One a) = One (One a) +initsDigit (Two a b) = Two (One a) (Two a b) +initsDigit (Three a b c) = Three (One a) (Two a b) (Three a b c) +initsDigit (Four a b c d) = Four (One a) (Two a b) (Three a b c) (Four a b c d) + +tailsNode :: Node a -> Node (Digit a) +tailsNode (Node2 s a b) = Node2 s (Two a b) (One b) +tailsNode (Node3 s a b c) = Node3 s (Three a b c) (Two b c) (One c) + +initsNode :: Node a -> Node (Digit a) +initsNode (Node2 s a b) = Node2 s (One a) (Two a b) +initsNode (Node3 s a b c) = Node3 s (One a) (Two a b) (Three a b c) + +{-# SPECIALIZE tailsTree :: (FingerTree (Elem a) -> Elem b) -> FingerTree (Elem a) -> FingerTree (Elem b) #-} +{-# SPECIALIZE tailsTree :: (FingerTree (Node a) -> Node b) -> FingerTree (Node a) -> FingerTree (Node b) #-} +-- | Given a function to apply to tails of a tree, applies that function +-- to every tail of the specified tree. +tailsTree :: (Sized a, Sized b) => (FingerTree a -> b) -> FingerTree a -> FingerTree b +tailsTree _ Empty = Empty +tailsTree f (Single x) = Single (f (Single x)) +tailsTree f (Deep n pr m sf) = + Deep n (fmap (\ pr' -> f (deep pr' m sf)) (tailsDigit pr)) + (tailsTree f' m) + (fmap (f . digitToTree) (tailsDigit sf)) + where + f' ms = let Just2 node m' = viewLTree ms in + fmap (\ pr' -> f (deep pr' m' sf)) (tailsNode node) + +{-# SPECIALIZE initsTree :: (FingerTree (Elem a) -> Elem b) -> FingerTree (Elem a) -> FingerTree (Elem b) #-} +{-# SPECIALIZE initsTree :: (FingerTree (Node a) -> Node b) -> FingerTree (Node a) -> FingerTree (Node b) #-} +-- | Given a function to apply to inits of a tree, applies that function +-- to every init of the specified tree. +initsTree :: (Sized a, Sized b) => (FingerTree a -> b) -> FingerTree a -> FingerTree b +initsTree _ Empty = Empty +initsTree f (Single x) = Single (f (Single x)) +initsTree f (Deep n pr m sf) = + Deep n (fmap (f . digitToTree) (initsDigit pr)) + (initsTree f' m) + (fmap (f . deep pr m) (initsDigit sf)) + where + f' ms = let Just2 m' node = viewRTree ms in + fmap (\ sf' -> f (deep pr m' sf')) (initsNode node) + +{-# INLINE foldlWithIndex #-} +-- | 'foldlWithIndex' is a version of 'foldl' that also provides access +-- to the index of each element. +foldlWithIndex :: (b -> Int -> a -> b) -> b -> Seq a -> b +foldlWithIndex f z xs = foldl (\ g x i -> i `seq` f (g (i - 1)) i x) (const z) xs (length xs - 1) + +{-# INLINE foldrWithIndex #-} +-- | 'foldrWithIndex' is a version of 'foldr' that also provides access +-- to the index of each element. +foldrWithIndex :: (Int -> a -> b -> b) -> b -> Seq a -> b +foldrWithIndex f z xs = foldr (\ x g i -> i `seq` f i x (g (i+1))) (const z) xs 0 + +{-# INLINE listToMaybe' #-} +-- 'listToMaybe\'' is a good consumer version of 'listToMaybe'. +listToMaybe' :: [a] -> Maybe a +listToMaybe' = foldr (\ x _ -> Just x) Nothing + +-- | /O(i)/ where /i/ is the prefix length. 'takeWhileL', applied +-- to a predicate @p@ and a sequence @xs@, returns the longest prefix +-- (possibly empty) of @xs@ of elements that satisfy @p@. +takeWhileL :: (a -> Bool) -> Seq a -> Seq a +takeWhileL p = fst . spanl p + +-- | /O(i)/ where /i/ is the suffix length. 'takeWhileR', applied +-- to a predicate @p@ and a sequence @xs@, returns the longest suffix +-- (possibly empty) of @xs@ of elements that satisfy @p@. +-- +-- @'takeWhileR' p xs@ is equivalent to @'reverse' ('takeWhileL' p ('reverse' xs))@. +takeWhileR :: (a -> Bool) -> Seq a -> Seq a +takeWhileR p = fst . spanr p + +-- | /O(i)/ where /i/ is the prefix length. @'dropWhileL' p xs@ returns +-- the suffix remaining after @'takeWhileL' p xs@. +dropWhileL :: (a -> Bool) -> Seq a -> Seq a +dropWhileL p = snd . spanl p + +-- | /O(i)/ where /i/ is the suffix length. @'dropWhileR' p xs@ returns +-- the prefix remaining after @'takeWhileR' p xs@. +-- +-- @'dropWhileR' p xs@ is equivalent to @'reverse' ('dropWhileL' p ('reverse' xs))@. +dropWhileR :: (a -> Bool) -> Seq a -> Seq a +dropWhileR p = snd . spanr p + +-- | /O(i)/ where /i/ is the prefix length. 'spanl', applied to +-- a predicate @p@ and a sequence @xs@, returns a pair whose first +-- element is the longest prefix (possibly empty) of @xs@ of elements that +-- satisfy @p@ and the second element is the remainder of the sequence. +spanl :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +spanl p = breakl (not . p) + +-- | /O(i)/ where /i/ is the suffix length. 'spanr', applied to a +-- predicate @p@ and a sequence @xs@, returns a pair whose /first/ element +-- is the longest /suffix/ (possibly empty) of @xs@ of elements that +-- satisfy @p@ and the second element is the remainder of the sequence. +spanr :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +spanr p = breakr (not . p) + +{-# INLINE breakl #-} +-- | /O(i)/ where /i/ is the breakpoint index. 'breakl', applied to a +-- predicate @p@ and a sequence @xs@, returns a pair whose first element +-- is the longest prefix (possibly empty) of @xs@ of elements that +-- /do not satisfy/ @p@ and the second element is the remainder of +-- the sequence. +-- +-- @'breakl' p@ is equivalent to @'spanl' (not . p)@. +breakl :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +breakl p xs = foldr (\ i _ -> splitAt i xs) (xs, empty) (findIndicesL p xs) + +{-# INLINE breakr #-} +-- | @'breakr' p@ is equivalent to @'spanr' (not . p)@. +breakr :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +breakr p xs = foldr (\ i _ -> flipPair (splitAt (i + 1) xs)) (xs, empty) (findIndicesR p xs) + where flipPair (x, y) = (y, x) + +-- | /O(n)/. The 'partition' function takes a predicate @p@ and a +-- sequence @xs@ and returns sequences of those elements which do and +-- do not satisfy the predicate. +partition :: (a -> Bool) -> Seq a -> (Seq a, Seq a) +partition p = foldl part (empty, empty) + where + part (xs, ys) x + | p x = (xs |> x, ys) + | otherwise = (xs, ys |> x) + +-- | /O(n)/. The 'filter' function takes a predicate @p@ and a sequence +-- @xs@ and returns a sequence of those elements which satisfy the +-- predicate. +filter :: (a -> Bool) -> Seq a -> Seq a +filter p = foldl (\ xs x -> if p x then xs |> x else xs) empty + +-- Indexing sequences + +-- | 'elemIndexL' finds the leftmost index of the specified element, +-- if it is present, and otherwise 'Nothing'. +elemIndexL :: Eq a => a -> Seq a -> Maybe Int +elemIndexL x = findIndexL (x ==) + +-- | 'elemIndexR' finds the rightmost index of the specified element, +-- if it is present, and otherwise 'Nothing'. +elemIndexR :: Eq a => a -> Seq a -> Maybe Int +elemIndexR x = findIndexR (x ==) + +-- | 'elemIndicesL' finds the indices of the specified element, from +-- left to right (i.e. in ascending order). +elemIndicesL :: Eq a => a -> Seq a -> [Int] +elemIndicesL x = findIndicesL (x ==) + +-- | 'elemIndicesR' finds the indices of the specified element, from +-- right to left (i.e. in descending order). +elemIndicesR :: Eq a => a -> Seq a -> [Int] +elemIndicesR x = findIndicesR (x ==) + +-- | @'findIndexL' p xs@ finds the index of the leftmost element that +-- satisfies @p@, if any exist. +findIndexL :: (a -> Bool) -> Seq a -> Maybe Int +findIndexL p = listToMaybe' . findIndicesL p + +-- | @'findIndexR' p xs@ finds the index of the rightmost element that +-- satisfies @p@, if any exist. +findIndexR :: (a -> Bool) -> Seq a -> Maybe Int +findIndexR p = listToMaybe' . findIndicesR p + +{-# INLINE findIndicesL #-} +-- | @'findIndicesL' p@ finds all indices of elements that satisfy @p@, +-- in ascending order. +findIndicesL :: (a -> Bool) -> Seq a -> [Int] +#if __GLASGOW_HASKELL__ +findIndicesL p xs = build (\ c n -> let g i x z = if p x then c i z else z in + foldrWithIndex g n xs) +#else +findIndicesL p xs = foldrWithIndex g [] xs + where g i x is = if p x then i:is else is +#endif + +{-# INLINE findIndicesR #-} +-- | @'findIndicesR' p@ finds all indices of elements that satisfy @p@, +-- in descending order. +findIndicesR :: (a -> Bool) -> Seq a -> [Int] +#if __GLASGOW_HASKELL__ +findIndicesR p xs = build (\ c n -> + let g z i x = if p x then c i z else z in foldlWithIndex g n xs) +#else +findIndicesR p xs = foldlWithIndex g [] xs + where g is i x = if p x then i:is else is +#endif + +------------------------------------------------------------------------ +-- Lists +------------------------------------------------------------------------ + +-- | /O(n)/. Create a sequence from a finite list of elements. +-- There is a function 'toList' in the opposite direction for all +-- instances of the 'Foldable' class, including 'Seq'. +fromList :: [a] -> Seq a +fromList = Data.List.foldl' (|>) empty + +------------------------------------------------------------------------ +-- Reverse +------------------------------------------------------------------------ + +-- | /O(n)/. The reverse of a sequence. +reverse :: Seq a -> Seq a +reverse (Seq xs) = Seq (reverseTree id xs) + +reverseTree :: (a -> a) -> FingerTree a -> FingerTree a +reverseTree _ Empty = Empty +reverseTree f (Single x) = Single (f x) +reverseTree f (Deep s pr m sf) = + Deep s (reverseDigit f sf) + (reverseTree (reverseNode f) m) + (reverseDigit f pr) + +{-# INLINE reverseDigit #-} +reverseDigit :: (a -> a) -> Digit a -> Digit a +reverseDigit f (One a) = One (f a) +reverseDigit f (Two a b) = Two (f b) (f a) +reverseDigit f (Three a b c) = Three (f c) (f b) (f a) +reverseDigit f (Four a b c d) = Four (f d) (f c) (f b) (f a) + +reverseNode :: (a -> a) -> Node a -> Node a +reverseNode f (Node2 s a b) = Node2 s (f b) (f a) +reverseNode f (Node3 s a b c) = Node3 s (f c) (f b) (f a) + +------------------------------------------------------------------------ +-- Zipping +------------------------------------------------------------------------ + +-- | /O(min(n1,n2))/. 'zip' takes two sequences and returns a sequence +-- of corresponding pairs. If one input is short, excess elements are +-- discarded from the right end of the longer sequence. +zip :: Seq a -> Seq b -> Seq (a, b) +zip = zipWith (,) + +-- | /O(min(n1,n2))/. 'zipWith' generalizes 'zip' by zipping with the +-- function given as the first argument, instead of a tupling function. +-- For example, @zipWith (+)@ is applied to two sequences to take the +-- sequence of corresponding sums. +zipWith :: (a -> b -> c) -> Seq a -> Seq b -> Seq c +zipWith f xs ys + | length xs <= length ys = zipWith' f xs ys + | otherwise = zipWith' (flip f) ys xs + +-- like 'zipWith', but assumes length xs <= length ys +zipWith' :: (a -> b -> c) -> Seq a -> Seq b -> Seq c +zipWith' f xs ys = snd (mapAccumL k ys xs) + where + k kys x = case viewl kys of + (z :< zs) -> (zs, f x z) + EmptyL -> error "zipWith': unexpected EmptyL" + +-- | /O(min(n1,n2,n3))/. 'zip3' takes three sequences and returns a +-- sequence of triples, analogous to 'zip'. +zip3 :: Seq a -> Seq b -> Seq c -> Seq (a,b,c) +zip3 = zipWith3 (,,) + +-- | /O(min(n1,n2,n3))/. 'zipWith3' takes a function which combines +-- three elements, as well as three sequences and returns a sequence of +-- their point-wise combinations, analogous to 'zipWith'. +zipWith3 :: (a -> b -> c -> d) -> Seq a -> Seq b -> Seq c -> Seq d +zipWith3 f s1 s2 s3 = zipWith ($) (zipWith f s1 s2) s3 + +-- | /O(min(n1,n2,n3,n4))/. 'zip4' takes four sequences and returns a +-- sequence of quadruples, analogous to 'zip'. +zip4 :: Seq a -> Seq b -> Seq c -> Seq d -> Seq (a,b,c,d) +zip4 = zipWith4 (,,,) + +-- | /O(min(n1,n2,n3,n4))/. 'zipWith4' takes a function which combines +-- four elements, as well as four sequences and returns a sequence of +-- their point-wise combinations, analogous to 'zipWith'. +zipWith4 :: (a -> b -> c -> d -> e) -> Seq a -> Seq b -> Seq c -> Seq d -> Seq e +zipWith4 f s1 s2 s3 s4 = zipWith ($) (zipWith ($) (zipWith f s1 s2) s3) s4 + +------------------------------------------------------------------------ +-- Sorting +-- +-- sort and sortBy are implemented by simple deforestations of +-- \ xs -> fromList2 (length xs) . Data.List.sortBy cmp . toList +-- which does not get deforested automatically, it would appear. +-- +-- Unstable sorting is performed by a heap sort implementation based on +-- pairing heaps. Because the internal structure of sequences is quite +-- varied, it is difficult to get blocks of elements of roughly the same +-- length, which would improve merge sort performance. Pairing heaps, +-- on the other hand, are relatively resistant to the effects of merging +-- heaps of wildly different sizes, as guaranteed by its amortized +-- constant-time merge operation. Moreover, extensive use of SpecConstr +-- transformations can be done on pairing heaps, especially when we're +-- only constructing them to immediately be unrolled. +-- +-- On purely random sequences of length 50000, with no RTS options, +-- I get the following statistics, in which heapsort is about 42.5% +-- faster: (all comparisons done with -O2) +-- +-- Times (ms) min mean +/-sd median max +-- to/from list: 103.802 108.572 7.487 106.436 143.339 +-- unstable heapsort: 60.686 62.968 4.275 61.187 79.151 +-- +-- Heapsort, it would seem, is less of a memory hog than Data.List.sortBy. +-- The gap is narrowed when more memory is available, but heapsort still +-- wins, 15% faster, with +RTS -H128m: +-- +-- Times (ms) min mean +/-sd median max +-- to/from list: 42.692 45.074 2.596 44.600 56.601 +-- unstable heapsort: 37.100 38.344 3.043 37.715 55.526 +-- +-- In addition, on strictly increasing sequences the gap is even wider +-- than normal; heapsort is 68.5% faster with no RTS options: +-- Times (ms) min mean +/-sd median max +-- to/from list: 52.236 53.574 1.987 53.034 62.098 +-- unstable heapsort: 16.433 16.919 0.931 16.681 21.622 +-- +-- This may be attributed to the elegant nature of the pairing heap. +-- +-- wasserman.louis@gmail.com, 7/20/09 +------------------------------------------------------------------------ + +-- | /O(n log n)/. 'sort' sorts the specified 'Seq' by the natural +-- ordering of its elements. The sort is stable. +-- If stability is not required, 'unstableSort' can be considerably +-- faster, and in particular uses less memory. +sort :: Ord a => Seq a -> Seq a +sort = sortBy compare + +-- | /O(n log n)/. 'sortBy' sorts the specified 'Seq' according to the +-- specified comparator. The sort is stable. +-- If stability is not required, 'unstableSortBy' can be considerably +-- faster, and in particular uses less memory. +sortBy :: (a -> a -> Ordering) -> Seq a -> Seq a +sortBy cmp xs = fromList2 (length xs) (Data.List.sortBy cmp (toList xs)) + +-- | /O(n log n)/. 'unstableSort' sorts the specified 'Seq' by +-- the natural ordering of its elements, but the sort is not stable. +-- This algorithm is frequently faster and uses less memory than 'sort', +-- and performs extremely well -- frequently twice as fast as 'sort' -- +-- when the sequence is already nearly sorted. +unstableSort :: Ord a => Seq a -> Seq a +unstableSort = unstableSortBy compare + +-- | /O(n log n)/. A generalization of 'unstableSort', 'unstableSortBy' +-- takes an arbitrary comparator and sorts the specified sequence. +-- The sort is not stable. This algorithm is frequently faster and +-- uses less memory than 'sortBy', and performs extremely well -- +-- frequently twice as fast as 'sortBy' -- when the sequence is already +-- nearly sorted. +unstableSortBy :: (a -> a -> Ordering) -> Seq a -> Seq a +unstableSortBy cmp (Seq xs) = + fromList2 (size xs) $ maybe [] (unrollPQ cmp) $ + toPQ cmp (\ (Elem x) -> PQueue x Nil) xs + +-- | fromList2, given a list and its length, constructs a completely +-- balanced Seq whose elements are that list using the applicativeTree +-- generalization. +fromList2 :: Int -> [a] -> Seq a +fromList2 n = execState (replicateA n (State ht)) + where + ht (x:xs) = (xs, x) + ht [] = error "fromList2: short list" + +-- | A 'PQueue' is a simple pairing heap. +data PQueue e = PQueue e (PQL e) +data PQL e = Nil | {-# UNPACK #-} !(PQueue e) :& PQL e + +infixr 8 :& + +#if TESTING + +instance Functor PQueue where + fmap f (PQueue x ts) = PQueue (f x) (fmap f ts) + +instance Functor PQL where + fmap f (q :& qs) = fmap f q :& fmap f qs + fmap _ Nil = Nil + +instance Show e => Show (PQueue e) where + show = unlines . draw . fmap show + +-- borrowed wholesale from Data.Tree, as Data.Tree actually depends +-- on Data.Sequence +draw :: PQueue String -> [String] +draw (PQueue x ts0) = x : drawSubTrees ts0 + where + drawSubTrees Nil = [] + drawSubTrees (t :& Nil) = + "|" : shift "`- " " " (draw t) + drawSubTrees (t :& ts) = + "|" : shift "+- " "| " (draw t) ++ drawSubTrees ts + + shift first other = Data.List.zipWith (++) (first : repeat other) +#endif + +-- | 'unrollPQ', given a comparator function, unrolls a 'PQueue' into +-- a sorted list. +unrollPQ :: (e -> e -> Ordering) -> PQueue e -> [e] +unrollPQ cmp = unrollPQ' + where + {-# INLINE unrollPQ' #-} + unrollPQ' (PQueue x ts) = x:mergePQs0 ts + (<>) = mergePQ cmp + mergePQs0 Nil = [] + mergePQs0 (t :& Nil) = unrollPQ' t + mergePQs0 (t1 :& t2 :& ts) = mergePQs (t1 <> t2) ts + mergePQs t ts = t `seq` case ts of + Nil -> unrollPQ' t + t1 :& Nil -> unrollPQ' (t <> t1) + t1 :& t2 :& ts' -> mergePQs (t <> (t1 <> t2)) ts' + +-- | 'toPQ', given an ordering function and a mechanism for queueifying +-- elements, converts a 'FingerTree' to a 'PQueue'. +toPQ :: (e -> e -> Ordering) -> (a -> PQueue e) -> FingerTree a -> Maybe (PQueue e) +toPQ _ _ Empty = Nothing +toPQ _ f (Single x) = Just (f x) +toPQ cmp f (Deep _ pr m sf) = Just (maybe (pr' <> sf') ((pr' <> sf') <>) (toPQ cmp fNode m)) + where + fDigit digit = case fmap f digit of + One a -> a + Two a b -> a <> b + Three a b c -> a <> b <> c + Four a b c d -> (a <> b) <> (c <> d) + (<>) = mergePQ cmp + fNode = fDigit . nodeToDigit + pr' = fDigit pr + sf' = fDigit sf + +-- | 'mergePQ' merges two 'PQueue's. +mergePQ :: (a -> a -> Ordering) -> PQueue a -> PQueue a -> PQueue a +mergePQ cmp q1@(PQueue x1 ts1) q2@(PQueue x2 ts2) + | cmp x1 x2 == GT = PQueue x2 (q1 :& ts2) + | otherwise = PQueue x1 (q2 :& ts1) diff --git a/Data/Set.hs b/Data/Set.hs new file mode 100644 --- /dev/null +++ b/Data/Set.hs @@ -0,0 +1,151 @@ +{-# 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 + , splitRoot + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , deleteAt + + -- * 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 + , link + , 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,1587 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +#if __GLASGOW_HASKELL__ >= 708 +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE TypeFamilies #-} +#endif +----------------------------------------------------------------------------- +-- | +-- 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 + , splitRoot + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , deleteAt + + -- * 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 + , link + , merge + ) where + +import Prelude hiding (filter,foldl,foldr,null,map) +import qualified Data.List as List +import Data.Bits (shiftL, shiftR) +import Data.Monoid (Monoid(..)) +import qualified Data.Foldable as Foldable +import Data.Typeable +import Control.DeepSeq (NFData(rnf)) + +import Data.StrictPair + +#if __GLASGOW_HASKELL__ +import GHC.Exts ( build ) +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts as GHCExts +#endif +import Text.Read +import Data.Data +#endif + +-- 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 + +{-------------------------------------------------------------------- + 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 + +#if __GLASGOW_HASKELL__ >= 708 +type role Set nominal +#endif + +instance Ord a => Monoid (Set a) where + mempty = empty + mappend = union + mconcat = unions + +instance Foldable.Foldable Set where + fold t = go t + where go Tip = mempty + go (Bin 1 k _ _) = k + go (Bin _ k l r) = go l `mappend` (k `mappend` go r) + {-# INLINABLE fold #-} + foldr = foldr + {-# INLINE foldr #-} + foldl = foldl + {-# INLINE foldl #-} + foldMap f t = go t + where go Tip = mempty + go (Bin 1 k _ _) = f k + go (Bin _ k l r) = go l `mappend` (f k `mappend` go r) + {-# INLINE foldMap #-} + +#if __GLASGOW_HASKELL__ + +{-------------------------------------------------------------------- + A Data instance +--------------------------------------------------------------------} + +-- This instance preserves data abstraction at the cost of inefficiency. +-- We provide limited reflection services for the sake of data abstraction. + +instance (Data a, Ord a) => Data (Set a) where + gfoldl f z set = z fromList `f` (toList set) + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = setDataType + dataCast1 f = gcast1 f + +fromListConstr :: Constr +fromListConstr = mkConstr setDataType "fromList" [] Prefix + +setDataType :: DataType +setDataType = mkDataType "Data.Set.Base.Set" [fromListConstr] + +#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 + +infix 4 `member` + +-- | /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 + +infix 4 `notMember` + +-- | /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. Returns an empty set if the set is empty. +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. Returns an empty set if the set is empty. +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. +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 + +infixl 5 `union` + +hedgeUnion :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a +hedgeUnion _ _ t1 Tip = t1 +hedgeUnion blo bhi Tip (Bin _ x l r) = link x (filterGt blo l) (filterLt bhi r) +hedgeUnion _ _ t1 (Bin _ x Tip Tip) = insertR x t1 -- According to benchmarks, this special case increases + -- performance up to 30%. It does not help in difference or intersection. +hedgeUnion blo bhi (Bin _ x l r) t2 = link x (hedgeUnion blo bmi l (trim blo bmi t2)) + (hedgeUnion bmi bhi r (trim bmi bhi t2)) + where bmi = JustS x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE hedgeUnion #-} +#endif + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} +-- | /O(n+m)/. Difference of two sets. +-- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/. +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 = link x (filterGt blo l) (filterLt bhi r) +hedgeDiff blo bhi t (Bin _ x l r) = merge (hedgeDiff blo bmi (trim blo bmi t) l) + (hedgeDiff bmi bhi (trim bmi bhi t) r) + where bmi = JustS x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE hedgeDiff #-} +#endif + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} +-- | /O(n+m)/. The intersection of two sets. The implementation uses an +-- efficient /hedge/ algorithm comparable with /hedge-union/. Elements of the +-- result come from the first set, so for example +-- +-- > 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 + +infixl 5 `intersection` + +hedgeInt :: Ord a => MaybeS a -> MaybeS a -> Set a -> Set a -> Set a +hedgeInt _ _ _ Tip = Tip +hedgeInt _ _ Tip _ = Tip +hedgeInt blo bhi (Bin _ x l r) t2 = let l' = hedgeInt blo bmi l (trim blo bmi t2) + r' = hedgeInt bmi bhi r (trim bmi bhi t2) + in if x `member` t2 then link x l' r' else merge l' r' + where bmi = JustS x +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE hedgeInt #-} +#endif + +{-------------------------------------------------------------------- + Filter and partition +--------------------------------------------------------------------} +-- | /O(n)/. Filter all elements that satisfy the predicate. +filter :: (a -> Bool) -> Set a -> Set a +filter _ Tip = Tip +filter p (Bin _ x l r) + | p x = link x (filter p l) (filter p r) + | otherwise = merge (filter p l) (filter p r) + +-- | /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 p0 t0 = toPair $ go p0 t0 + where + go _ Tip = (Tip :*: Tip) + go p (Bin _ x l r) = case (go p l, go p r) of + ((l1 :*: l2), (r1 :*: r2)) + | p x -> link x l1 r1 :*: merge l2 r2 + | otherwise -> merge l1 r1 :*: link 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 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 +--------------------------------------------------------------------} +#if __GLASGOW_HASKELL__ >= 708 +instance (Ord a) => GHCExts.IsList (Set a) where + type Item (Set a) = a + fromList = fromList + toList = toList +#endif + +-- | /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. +-- +-- If the elemens are ordered, linear-time implementation is used, +-- with the performance equal to 'fromDistinctAscList'. + +-- For some reason, when 'singleton' is used in fromList or in +-- create, it is not inlined, so we inline it manually. +fromList :: Ord a => [a] -> Set a +fromList [] = Tip +fromList [x] = Bin 1 x Tip Tip +fromList (x0 : xs0) | not_ordered x0 xs0 = fromList' (Bin 1 x0 Tip Tip) xs0 + | otherwise = go (1::Int) (Bin 1 x0 Tip Tip) xs0 + where + not_ordered _ [] = False + not_ordered x (y : _) = x >= y + {-# INLINE not_ordered #-} + + fromList' t0 xs = foldlStrict ins t0 xs + where ins t x = insert x t + + STRICT_1_OF_3(go) + go _ t [] = t + go _ t [x] = insertMax x t + go s l xs@(x : xss) | not_ordered x xss = fromList' l xs + | otherwise = case create s xss of + (r, ys, []) -> go (s `shiftL` 1) (link x l r) ys + (r, _, ys) -> fromList' (link x l r) ys + + -- The create is returning a triple (tree, xs, ys). Both xs and ys + -- represent not yet processed elements and only one of them can be nonempty. + -- If ys is nonempty, the keys in ys are not ordered with respect to tree + -- and must be inserted using fromList'. Otherwise the keys have been + -- ordered so far. + STRICT_1_OF_2(create) + create _ [] = (Tip, [], []) + create s xs@(x : xss) + | s == 1 = if not_ordered x xss then (Bin 1 x Tip Tip, [], xss) + else (Bin 1 x Tip Tip, xss, []) + | otherwise = case create (s `shiftR` 1) xs of + res@(_, [], _) -> res + (l, [y], zs) -> (insertMax y l, [], zs) + (l, ys@(y:yss), _) | not_ordered y yss -> (l, [], ys) + | otherwise -> case create (s `shiftR` 1) yss of + (r, zs, ws) -> (link y l r, zs, ws) +#if __GLASGOW_HASKELL__ >= 700 +{-# 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./ + +-- For some reason, when 'singleton' is used in fromDistinctAscList or in +-- create, it is not inlined, so we inline it manually. +fromDistinctAscList :: [a] -> Set a +fromDistinctAscList [] = Tip +fromDistinctAscList (x0 : xs0) = go (1::Int) (Bin 1 x0 Tip Tip) xs0 + where + STRICT_1_OF_3(go) + go _ t [] = t + go s l (x : xs) = case create s xs of + (r, ys) -> go (s `shiftL` 1) (link x l r) ys + + STRICT_1_OF_2(create) + create _ [] = (Tip, []) + create s xs@(x : xs') + | s == 1 = (Bin 1 x Tip Tip, xs') + | otherwise = case create (s `shiftR` 1) xs of + res@(_, []) -> res + (l, y:ys) -> case create (s `shiftR` 1) ys of + (r, zs) -> (link y l r, zs) + +{-------------------------------------------------------------------- + Eq converts the set to a list. In a lazy setting, this + actually seems one of the faster methods to compare two trees + 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 -> link x (filter' b' l) r + EQ -> r + GT -> filter' b' r +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE filterGt #-} +#endif + +filterLt :: Ord a => MaybeS a -> Set a -> Set a +filterLt NothingS t = t +filterLt (JustS b) t = filter' b t + where filter' _ Tip = Tip + filter' b' (Bin _ x l r) = + case compare x b' of LT -> link x l (filter' b' r) + EQ -> l + GT -> filter' b' l +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE filterLt #-} +#endif + +{-------------------------------------------------------------------- + Split +--------------------------------------------------------------------} +-- | /O(log n)/. The expression (@'split' x set@) is a pair @(set1,set2)@ +-- where @set1@ comprises the elements of @set@ less than @x@ and @set2@ +-- comprises the elements of @set@ greater than @x@. +split :: Ord a => a -> Set a -> (Set a,Set a) +split x0 t0 = toPair $ go x0 t0 + where + go _ Tip = (Tip :*: Tip) + go x (Bin _ y l r) + = case compare x y of + LT -> let (lt :*: gt) = go x l in (lt :*: link y gt r) + GT -> let (lt :*: gt) = go x r in (link y l lt :*: gt) + 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 + gt' = link y gt r + in gt' `seq` (lt, found, gt') + GT -> let (lt, found, gt) = splitMember x r + lt' = link y l lt + in lt' `seq` (lt', found, gt) + EQ -> (l, True, r) +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE splitMember #-} +#endif + +{-------------------------------------------------------------------- + Indexing +--------------------------------------------------------------------} + +-- | /O(log n)/. Return the /index/ of an element, which is its zero-based +-- index in the sorted sequence of elements. The index is a number from /0/ up +-- to, but not including, the 'size' of the set. Calls 'error' when the element +-- is not a 'member' of the set. +-- +-- > findIndex 2 (fromList [5,3]) Error: element is not in the set +-- > findIndex 3 (fromList [5,3]) == 0 +-- > findIndex 5 (fromList [5,3]) == 1 +-- > findIndex 6 (fromList [5,3]) Error: element is not in the set + +-- See Note: Type of local 'go' function +findIndex :: Ord a => a -> Set a -> Int +findIndex = go 0 + where + go :: Ord a => Int -> a -> Set a -> Int + STRICT_1_OF_3(go) + STRICT_2_OF_3(go) + go _ _ Tip = error "Set.findIndex: element is not in the set" + go idx x (Bin _ kx l r) = case compare x kx of + LT -> go idx x l + GT -> go (idx + size l + 1) x r + EQ -> idx + size l +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE findIndex #-} +#endif + +-- | /O(log n)/. Lookup the /index/ of an element, which is its zero-based index in +-- the sorted sequence of elements. The index is a number from /0/ up to, but not +-- including, the 'size' of the set. +-- +-- > isJust (lookupIndex 2 (fromList [5,3])) == False +-- > fromJust (lookupIndex 3 (fromList [5,3])) == 0 +-- > fromJust (lookupIndex 5 (fromList [5,3])) == 1 +-- > isJust (lookupIndex 6 (fromList [5,3])) == False + +-- See Note: Type of local 'go' function +lookupIndex :: Ord a => a -> Set a -> Maybe Int +lookupIndex = go 0 + where + go :: Ord a => Int -> a -> Set a -> Maybe Int + STRICT_1_OF_3(go) + STRICT_2_OF_3(go) + go _ _ Tip = Nothing + go idx x (Bin _ kx l r) = case compare x kx of + LT -> go idx x l + GT -> go (idx + size l + 1) x r + EQ -> Just $! idx + size l +#if __GLASGOW_HASKELL__ >= 700 +{-# INLINABLE lookupIndex #-} +#endif + +-- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based +-- index in the sorted sequence of elements. If the /index/ is out of range (less +-- than zero, greater or equal to 'size' of the set), 'error' is called. +-- +-- > elemAt 0 (fromList [5,3]) == 3 +-- > elemAt 1 (fromList [5,3]) == 5 +-- > elemAt 2 (fromList [5,3]) Error: index out of range + +elemAt :: Int -> Set a -> a +STRICT_1_OF_2(elemAt) +elemAt _ Tip = error "Set.elemAt: index out of range" +elemAt i (Bin _ x l r) + = case compare i sizeL of + LT -> elemAt i l + GT -> elemAt (i-sizeL-1) r + EQ -> x + where + sizeL = size l + +-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in +-- the sorted sequence of elements. If the /index/ is out of range (less than zero, +-- greater or equal to 'size' of the set), 'error' is called. +-- +-- > deleteAt 0 (fromList [5,3]) == singleton 5 +-- > deleteAt 1 (fromList [5,3]) == singleton 3 +-- > deleteAt 2 (fromList [5,3]) Error: index out of range +-- > deleteAt (-1) (fromList [5,3]) Error: index out of range + +deleteAt :: Int -> Set a -> Set a +deleteAt i t = i `seq` + case t of + Tip -> error "Set.deleteAt: index out of range" + Bin _ x l r -> case compare i sizeL of + LT -> balanceR x (deleteAt i l) r + GT -> balanceL x l (deleteAt (i-sizeL-1) r) + EQ -> glue l r + where + sizeL = size l + + +{-------------------------------------------------------------------- + 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. + [link x l r] Restores balance and size. + + Furthermore, we can construct a new tree from two trees. Both operations + assume that all values in [l] < all values in [r] and that [l] and [r] + are valid: + [glue l r] Glues [l] and [r] together. Assumes that [l] and + [r] are already balanced with respect to each other. + [merge l r] Merges two trees and restores balance. + + Note: in contrast to Adam's paper, we use (<=) comparisons instead + of (<) comparisons in [link], [merge] and [balance]. + Quickcheck (on [difference]) showed that this was necessary in order + to maintain the invariants. It is quite unsatisfactory that I haven't + been able to find out why this is actually the case! Fortunately, it + doesn't hurt to be a bit more conservative. +--------------------------------------------------------------------} + +{-------------------------------------------------------------------- + Link +--------------------------------------------------------------------} +link :: a -> Set a -> Set a -> Set a +link x Tip r = insertMin x r +link x l Tip = insertMax x l +link x l@(Bin sizeL y ly ry) r@(Bin sizeR z lz rz) + | delta*sizeL < sizeR = balanceL z (link x l lz) rz + | delta*sizeR < sizeL = balanceR y ly (link 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 #-} + +-- | /O(1)/. Decompose a set into pieces based on the structure of the underlying +-- tree. This function is useful for consuming a set in parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that the pieces +-- returned will be in ascending order (all elements in the first subset less than all +-- elements in the second, and so on). +-- +-- Examples: +-- +-- > splitRoot (fromList [1..6]) == +-- > [fromList [1,2,3],fromList [4],fromList [5,6]] +-- +-- > splitRoot empty == [] +-- +-- Note that the current implementation does not return more than three subsets, +-- but you should not depend on this behaviour because it can change in the +-- future without notice. +splitRoot :: Set a -> [Set a] +splitRoot orig = + case orig of + Tip -> [] + Bin _ v l r -> [l, singleton v, r] +{-# INLINE splitRoot #-} + + +{-------------------------------------------------------------------- + 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,13 @@ +{-# LANGUAGE CPP #-} +#if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Trustworthy #-} +#endif +module Data.StrictPair (StrictPair(..), toPair) where + +-- | Same as regular Haskell pairs, but (x :*: _|_) = (_|_ :*: y) = +-- _|_ +data StrictPair a b = !a :*: !b + +toPair :: StrictPair a b -> (a, b) +toPair (x :*: y) = (x, y) +{-# INLINE toPair #-} diff --git a/Data/Tree.hs b/Data/Tree.hs new file mode 100644 --- /dev/null +++ b/Data/Tree.hs @@ -0,0 +1,170 @@ +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} +#endif +#if __GLASGOW_HASKELL__ >= 703 +{-# LANGUAGE Safe #-} +#endif +----------------------------------------------------------------------------- +-- | +-- 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 +-- +-- Multi-way trees (/aka/ rose trees) and forests. +-- +----------------------------------------------------------------------------- + +module Data.Tree( + Tree(..), Forest, + -- * Two-dimensional drawing + drawTree, drawForest, + -- * Extraction + flatten, levels, + -- * Building trees + unfoldTree, unfoldForest, + unfoldTreeM, unfoldForestM, + unfoldTreeM_BF, unfoldForestM_BF, + ) where + +import Control.Applicative (Applicative(..), (<$>)) +import Control.Monad (liftM) +import Data.Monoid (Monoid(..)) +import Data.Sequence (Seq, empty, singleton, (<|), (|>), fromList, + ViewL(..), ViewR(..), viewl, viewr) +import Data.Foldable (Foldable(foldMap), toList) +import Data.Traversable (Traversable(traverse)) +import Data.Typeable +import Control.DeepSeq (NFData(rnf)) + +#ifdef __GLASGOW_HASKELL__ +import Data.Data (Data) +#endif + +-- | Multi-way trees, also known as /rose trees/. +data Tree a = Node { + rootLabel :: a, -- ^ label value + subForest :: Forest a -- ^ zero or more child trees + } +#ifdef __GLASGOW_HASKELL__ + deriving (Eq, Read, Show, Data) +#else + deriving (Eq, Read, Show) +#endif +type Forest a = [Tree a] + +#include "Typeable.h" +INSTANCE_TYPEABLE1(Tree,treeTc,"Tree") + +instance Functor Tree where + fmap f (Node x ts) = Node (f x) (map (fmap f) ts) + +instance Applicative Tree where + pure x = Node x [] + Node f tfs <*> tx@(Node x txs) = + Node (f x) (map (f <$>) txs ++ map (<*> tx) tfs) + +instance Monad Tree where + return x = Node x [] + Node x ts >>= f = Node x' (ts' ++ map (>>= f) ts) + where Node x' ts' = f x + +instance Traversable Tree where + traverse f (Node x ts) = Node <$> f x <*> traverse (traverse f) ts + +instance Foldable Tree where + foldMap f (Node x ts) = f x `mappend` foldMap (foldMap f) ts + +instance NFData a => NFData (Tree a) where + rnf (Node x ts) = rnf x `seq` rnf ts + +-- | Neat 2-dimensional drawing of a tree. +drawTree :: Tree String -> String +drawTree = unlines . draw + +-- | Neat 2-dimensional drawing of a forest. +drawForest :: Forest String -> String +drawForest = unlines . map drawTree + +draw :: Tree String -> [String] +draw (Node x ts0) = x : drawSubTrees ts0 + where + drawSubTrees [] = [] + drawSubTrees [t] = + "|" : shift "`- " " " (draw t) + drawSubTrees (t:ts) = + "|" : shift "+- " "| " (draw t) ++ drawSubTrees ts + + shift first other = zipWith (++) (first : repeat other) + +-- | The elements of a tree in pre-order. +flatten :: Tree a -> [a] +flatten t = squish t [] + where squish (Node x ts) xs = x:Prelude.foldr squish xs ts + +-- | Lists of nodes at each level of the tree. +levels :: Tree a -> [[a]] +levels t = + map (map rootLabel) $ + takeWhile (not . null) $ + iterate (concatMap subForest) [t] + +-- | Build a tree from a seed value +unfoldTree :: (b -> (a, [b])) -> b -> Tree a +unfoldTree f b = let (a, bs) = f b in Node a (unfoldForest f bs) + +-- | Build a forest from a list of seed values +unfoldForest :: (b -> (a, [b])) -> [b] -> Forest a +unfoldForest f = map (unfoldTree f) + +-- | Monadic tree builder, in depth-first order +unfoldTreeM :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a) +unfoldTreeM f b = do + (a, bs) <- f b + ts <- unfoldForestM f bs + return (Node a ts) + +-- | Monadic forest builder, in depth-first order +#ifndef __NHC__ +unfoldForestM :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a) +#endif +unfoldForestM f = Prelude.mapM (unfoldTreeM f) + +-- | Monadic tree builder, in breadth-first order, +-- using an algorithm adapted from +-- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/, +-- by Chris Okasaki, /ICFP'00/. +unfoldTreeM_BF :: Monad m => (b -> m (a, [b])) -> b -> m (Tree a) +unfoldTreeM_BF f b = liftM getElement $ unfoldForestQ f (singleton b) + where + getElement xs = case viewl xs of + x :< _ -> x + EmptyL -> error "unfoldTreeM_BF" + +-- | Monadic forest builder, in breadth-first order, +-- using an algorithm adapted from +-- /Breadth-First Numbering: Lessons from a Small Exercise in Algorithm Design/, +-- by Chris Okasaki, /ICFP'00/. +unfoldForestM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m (Forest a) +unfoldForestM_BF f = liftM toList . unfoldForestQ f . fromList + +-- takes a sequence (queue) of seeds +-- produces a sequence (reversed queue) of trees of the same length +unfoldForestQ :: Monad m => (b -> m (a, [b])) -> Seq b -> m (Seq (Tree a)) +unfoldForestQ f aQ = case viewl aQ of + EmptyL -> return empty + a :< aQ' -> do + (b, as) <- f a + tQ <- unfoldForestQ f (Prelude.foldl (|>) aQ' as) + let (tQ', ts) = splitOnto [] as tQ + return (Node b ts <| tQ') + where + splitOnto :: [a'] -> [b'] -> Seq a' -> (Seq a', [a']) + splitOnto as [] q = (q, as) + splitOnto as (_:bs) q = case viewr q of + q' :> a -> splitOnto (a:as) bs q' + EmptyR -> error "unfoldForestQ" diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,31 @@ +The Glasgow Haskell Compiler License + +Copyright 2004, The University Court of the University of Glasgow. +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +- 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. + +THIS SOFTWARE IS PROVIDED BY THE UNIVERSITY COURT OF THE UNIVERSITY OF +GLASGOW AND THE CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, +INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +UNIVERSITY COURT OF THE UNIVERSITY OF GLASGOW OR THE CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY +OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +DAMAGE. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,6 @@ +module Main (main) where + +import Distribution.Simple + +main :: IO () +main = defaultMain diff --git a/benchmarks/IntMap.hs b/benchmarks/IntMap.hs new file mode 100644 --- /dev/null +++ b/benchmarks/IntMap.hs @@ -0,0 +1,97 @@ +{-# 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 + , bench "fromList" $ whnf M.fromList elems + , bench "fromAscList" $ whnf M.fromAscList elems + , bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems + ] + 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,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.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 + , bench "fromList" $ whnf S.fromList elems + , bench "fromAscList" $ whnf S.fromAscList elems + , bench "fromDistinctAscList" $ whnf S.fromDistinctAscList elems + ] + where + elems = [1..2^12] + elems_even = [2,4..2^12] + elems_odd = [1,3..2^12] + +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-$* $(BENCHMARK) -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,131 @@ +{-# 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 + , bench "split" $ whnf (M.split (bound `div` 2)) m + , bench "fromList" $ whnf M.fromList elems + , bench "fromList-desc" $ whnf M.fromList (reverse elems) + , bench "fromAscList" $ whnf M.fromAscList elems + , bench "fromDistinctAscList" $ whnf M.fromDistinctAscList elems + ] + where + bound = 2^12 + 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,53 @@ +{-# 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 + , bench "fromList" $ whnf S.fromList elems + , bench "fromList-desc" $ whnf S.fromList (reverse elems) + , bench "fromAscList" $ whnf S.fromAscList elems + , bench "fromDistinctAscList" $ whnf S.fromDistinctAscList elems + ] + where + elems = [1..2^12] + elems_even = [2,4..2^12] + elems_odd = [1,3..2^12] + +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/exposed-containers.cabal b/exposed-containers.cabal new file mode 100644 --- /dev/null +++ b/exposed-containers.cabal @@ -0,0 +1,246 @@ +name: exposed-containers +version: 0.5.5.1 +license: BSD3 +license-file: LICENSE +maintainer: fox@ucw.cz +bug-reports: https://github.com/haskell/containers/issues +synopsis: A distribution of the 'containers' package, with all modules exposed. +category: Data Structures +description: + The source package contains efficient general-purpose implementations of + various basic immutable container types. The declared cost of each + operation is either worst-case or amortized, but remains valid even if + structures are shared. + + Here we redistribute it, but with hidden modules exposed. +build-type: Simple +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/fmap/exposed-containers.git + +Library + build-depends: base >= 4.2 && < 5, array, deepseq >= 1.2 && < 1.4 + 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 + Data.BitUtil + Data.IntMap.Base + Data.IntSet.Base + Data.Map.Base + Data.Set.Base + Data.StrictPair + if !impl(nhc98) + exposed-modules: + Data.Graph + Data.Sequence + Data.Tree + + include-dirs: include + + if impl(ghc<7.0) + extensions: MagicHash, DeriveDataTypeable, StandaloneDeriving, Rank2Types + if impl(ghc >= 7.8) + extensions: RoleAnnotations + +------------------- +-- 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 deprecated-properties + hs-source-dirs: tests, . + main-is: deprecated-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 + +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 + +test-suite map-strictness-properties + hs-source-dirs: tests, . + main-is: MapStrictness.hs + type: exitcode-stdio-1.0 + + build-depends: + array, + base >= 4.2 && < 5, + ChasingBottoms, + deepseq >= 1.2 && < 1.4, + QuickCheck >= 2.4.0.1, + ghc-prim, + test-framework >= 0.3.3, + test-framework-quickcheck2 >= 0.2.9 + + ghc-options: -Wall + +test-suite intmap-strictness-properties + hs-source-dirs: tests, . + main-is: IntMapStrictness.hs + type: exitcode-stdio-1.0 + + build-depends: + array, + base >= 4.2 && < 5, + ChasingBottoms, + deepseq >= 1.2 && < 1.4, + QuickCheck >= 2.4.0.1, + ghc-prim, + test-framework >= 0.3.3, + test-framework-quickcheck2 >= 0.2.9 + + ghc-options: -Wall diff --git a/include/Typeable.h b/include/Typeable.h new file mode 100644 --- /dev/null +++ b/include/Typeable.h @@ -0,0 +1,65 @@ +{- -------------------------------------------------------------------------- +// Macros to help make Typeable instances. +// +// 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) +// -------------------------------------------------------------------------- +-} + +#ifndef TYPEABLE_H +#define TYPEABLE_H + +#ifdef __GLASGOW_HASKELL__ + +-- // For GHC, we can use DeriveDataTypeable + StandaloneDeriving to +-- // generate the instances. + +#define INSTANCE_TYPEABLE0(tycon,tcname,str) deriving instance Typeable tycon +#if __GLASGOW_HASKELL__ >= 707 +#define INSTANCE_TYPEABLE1(tycon,tcname,str) deriving instance Typeable tycon +#define INSTANCE_TYPEABLE2(tycon,tcname,str) deriving instance Typeable tycon +#define INSTANCE_TYPEABLE3(tycon,tcname,str) deriving instance Typeable tycon +#else +#define INSTANCE_TYPEABLE1(tycon,tcname,str) deriving instance Typeable1 tycon +#define INSTANCE_TYPEABLE2(tycon,tcname,str) deriving instance Typeable2 tycon +#define INSTANCE_TYPEABLE3(tycon,tcname,str) deriving instance Typeable3 tycon +#endif + +#else /* !__GLASGOW_HASKELL__ */ + +#define INSTANCE_TYPEABLE0(tycon,tcname,str) \ +tcname :: TyCon; \ +tcname = mkTyCon str; \ +instance Typeable tycon where { typeOf _ = mkTyConApp tcname [] } + +#define INSTANCE_TYPEABLE1(tycon,tcname,str) \ +tcname = mkTyCon str; \ +instance Typeable1 tycon where { typeOf1 _ = mkTyConApp tcname [] }; \ +instance Typeable a => Typeable (tycon a) where { typeOf = typeOfDefault } + +#define INSTANCE_TYPEABLE2(tycon,tcname,str) \ +tcname = mkTyCon str; \ +instance Typeable2 tycon where { typeOf2 _ = mkTyConApp tcname [] }; \ +instance Typeable a => Typeable1 (tycon a) where { \ + typeOf1 = typeOf1Default }; \ +instance (Typeable a, Typeable b) => Typeable (tycon a b) where { \ + typeOf = typeOfDefault } + +#define INSTANCE_TYPEABLE3(tycon,tcname,str) \ +tcname = mkTyCon str; \ +instance Typeable3 tycon where { typeOf3 _ = mkTyConApp tcname [] }; \ +instance Typeable a => Typeable2 (tycon a) where { \ + typeOf2 = typeOf2Default }; \ +instance (Typeable a, Typeable b) => Typeable1 (tycon a b) where { \ + typeOf1 = typeOf1Default }; \ +instance (Typeable a, Typeable b, Typeable c) => Typeable (tycon a b c) where { \ + typeOf = typeOfDefault } + +#endif /* !__GLASGOW_HASKELL__ */ + +#endif diff --git a/tests/IntMapStrictness.hs b/tests/IntMapStrictness.hs new file mode 100644 --- /dev/null +++ b/tests/IntMapStrictness.hs @@ -0,0 +1,127 @@ +{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Main (main) where + +import Test.ChasingBottoms.IsBottom +import Test.Framework (Test, defaultMain, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck (Arbitrary(arbitrary)) + +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as M + +instance Arbitrary v => Arbitrary (IntMap v) where + arbitrary = M.fromList `fmap` arbitrary + +instance Show (Int -> Int) where + show _ = "" + +instance Show (Int -> Int -> Int) where + show _ = "" + +instance Show (Int -> Int -> Int -> Int) where + show _ = "" + +------------------------------------------------------------------------ +-- * Properties + +------------------------------------------------------------------------ +-- ** Strict module + +pSingletonKeyStrict :: Int -> Bool +pSingletonKeyStrict v = isBottom $ M.singleton (bottom :: Int) v + +pSingletonValueStrict :: Int -> Bool +pSingletonValueStrict k = isBottom $ (M.singleton k (bottom :: Int)) + +pFindWithDefaultKeyStrict :: Int -> IntMap Int -> Bool +pFindWithDefaultKeyStrict def m = isBottom $ M.findWithDefault def bottom m + +pFindWithDefaultValueStrict :: Int -> IntMap Int -> Bool +pFindWithDefaultValueStrict k m = + M.member k m || (isBottom $ M.findWithDefault bottom k m) + +pAdjustKeyStrict :: (Int -> Int) -> IntMap Int -> Bool +pAdjustKeyStrict f m = isBottom $ M.adjust f bottom m + +pAdjustValueStrict :: Int -> IntMap Int -> Bool +pAdjustValueStrict k m + | k `M.member` m = isBottom $ M.adjust (const bottom) k m + | otherwise = case M.keys m of + [] -> True + (k':_) -> isBottom $ M.adjust (const bottom) k' m + +pInsertKeyStrict :: Int -> IntMap Int -> Bool +pInsertKeyStrict v m = isBottom $ M.insert bottom v m + +pInsertValueStrict :: Int -> IntMap Int -> Bool +pInsertValueStrict k m = isBottom $ M.insert k bottom m + +pInsertWithKeyStrict :: (Int -> Int -> Int) -> Int -> IntMap Int -> Bool +pInsertWithKeyStrict f v m = isBottom $ M.insertWith f bottom v m + +pInsertWithValueStrict :: (Int -> Int -> Int) -> Int -> Int -> IntMap Int + -> Bool +pInsertWithValueStrict f k v m + | M.member k m = (isBottom $ M.insertWith (const2 bottom) k v m) && + not (isBottom $ M.insertWith (const2 1) k bottom m) + | otherwise = isBottom $ M.insertWith f k bottom m + +pInsertLookupWithKeyKeyStrict :: (Int -> Int -> Int -> Int) -> Int -> IntMap Int + -> Bool +pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey f bottom v m + +pInsertLookupWithKeyValueStrict :: (Int -> Int -> Int -> Int) -> Int -> Int + -> IntMap Int -> Bool +pInsertLookupWithKeyValueStrict f k v m + | M.member k m = (isBottom $ M.insertLookupWithKey (const3 bottom) k v m) && + not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m) + | otherwise = isBottom $ M.insertLookupWithKey f k bottom m + +------------------------------------------------------------------------ +-- * Test list + +tests :: [Test] +tests = + [ + -- Basic interface + testGroup "IntMap.Strict" + [ testProperty "singleton is key-strict" pSingletonKeyStrict + , testProperty "singleton is value-strict" pSingletonValueStrict + , testProperty "member is key-strict" $ keyStrict M.member + , testProperty "lookup is key-strict" $ keyStrict M.lookup + , testProperty "findWithDefault is key-strict" pFindWithDefaultKeyStrict + , testProperty "findWithDefault is value-strict" pFindWithDefaultValueStrict + , testProperty "! is key-strict" $ keyStrict (flip (M.!)) + , testProperty "delete is key-strict" $ keyStrict M.delete + , testProperty "adjust is key-strict" pAdjustKeyStrict + , testProperty "adjust is value-strict" pAdjustValueStrict + , testProperty "insert is key-strict" pInsertKeyStrict + , testProperty "insert is value-strict" pInsertValueStrict + , testProperty "insertWith is key-strict" pInsertWithKeyStrict + , testProperty "insertWith is value-strict" pInsertWithValueStrict + , testProperty "insertLookupWithKey is key-strict" + pInsertLookupWithKeyKeyStrict + , testProperty "insertLookupWithKey is value-strict" + pInsertLookupWithKeyValueStrict + ] + ] + +------------------------------------------------------------------------ +-- * Test harness + +main :: IO () +main = defaultMain tests + +------------------------------------------------------------------------ +-- * Utilities + +keyStrict :: (Int -> IntMap Int -> a) -> IntMap Int -> Bool +keyStrict f m = isBottom $ f bottom m + +const2 :: a -> b -> c -> a +const2 x _ _ = x + +const3 :: a -> b -> c -> d -> a +const3 x _ _ _ = x 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/MapStrictness.hs b/tests/MapStrictness.hs new file mode 100644 --- /dev/null +++ b/tests/MapStrictness.hs @@ -0,0 +1,128 @@ +{-# LANGUAGE FlexibleInstances, GeneralizedNewtypeDeriving #-} +{-# OPTIONS_GHC -fno-warn-orphans #-} + +module Main (main) where + +import Test.ChasingBottoms.IsBottom +import Test.Framework (Test, defaultMain, testGroup) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck (Arbitrary(arbitrary)) + +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as M + +instance (Arbitrary k, Arbitrary v, Eq k, Ord k) => + Arbitrary (Map k v) where + arbitrary = M.fromList `fmap` arbitrary + +instance Show (Int -> Int) where + show _ = "" + +instance Show (Int -> Int -> Int) where + show _ = "" + +instance Show (Int -> Int -> Int -> Int) where + show _ = "" + +------------------------------------------------------------------------ +-- * Properties + +------------------------------------------------------------------------ +-- ** Strict module + +pSingletonKeyStrict :: Int -> Bool +pSingletonKeyStrict v = isBottom $ M.singleton (bottom :: Int) v + +pSingletonValueStrict :: Int -> Bool +pSingletonValueStrict k = isBottom $ (M.singleton k (bottom :: Int)) + +pFindWithDefaultKeyStrict :: Int -> Map Int Int -> Bool +pFindWithDefaultKeyStrict def m = isBottom $ M.findWithDefault def bottom m + +pFindWithDefaultValueStrict :: Int -> Map Int Int -> Bool +pFindWithDefaultValueStrict k m = + M.member k m || (isBottom $ M.findWithDefault bottom k m) + +pAdjustKeyStrict :: (Int -> Int) -> Map Int Int -> Bool +pAdjustKeyStrict f m = isBottom $ M.adjust f bottom m + +pAdjustValueStrict :: Int -> Map Int Int -> Bool +pAdjustValueStrict k m + | k `M.member` m = isBottom $ M.adjust (const bottom) k m + | otherwise = case M.keys m of + [] -> True + (k':_) -> isBottom $ M.adjust (const bottom) k' m + +pInsertKeyStrict :: Int -> Map Int Int -> Bool +pInsertKeyStrict v m = isBottom $ M.insert bottom v m + +pInsertValueStrict :: Int -> Map Int Int -> Bool +pInsertValueStrict k m = isBottom $ M.insert k bottom m + +pInsertWithKeyStrict :: (Int -> Int -> Int) -> Int -> Map Int Int -> Bool +pInsertWithKeyStrict f v m = isBottom $ M.insertWith f bottom v m + +pInsertWithValueStrict :: (Int -> Int -> Int) -> Int -> Int -> Map Int Int + -> Bool +pInsertWithValueStrict f k v m + | M.member k m = (isBottom $ M.insertWith (const2 bottom) k v m) && + not (isBottom $ M.insertWith (const2 1) k bottom m) + | otherwise = isBottom $ M.insertWith f k bottom m + +pInsertLookupWithKeyKeyStrict :: (Int -> Int -> Int -> Int) -> Int + -> Map Int Int -> Bool +pInsertLookupWithKeyKeyStrict f v m = isBottom $ M.insertLookupWithKey f bottom v m + +pInsertLookupWithKeyValueStrict :: (Int -> Int -> Int -> Int) -> Int -> Int + -> Map Int Int -> Bool +pInsertLookupWithKeyValueStrict f k v m + | M.member k m = (isBottom $ M.insertLookupWithKey (const3 bottom) k v m) && + not (isBottom $ M.insertLookupWithKey (const3 1) k bottom m) + | otherwise = isBottom $ M.insertLookupWithKey f k bottom m + +------------------------------------------------------------------------ +-- * Test list + +tests :: [Test] +tests = + [ + -- Basic interface + testGroup "Map.Strict" + [ testProperty "singleton is key-strict" pSingletonKeyStrict + , testProperty "singleton is value-strict" pSingletonValueStrict + , testProperty "member is key-strict" $ keyStrict M.member + , testProperty "lookup is key-strict" $ keyStrict M.lookup + , testProperty "findWithDefault is key-strict" pFindWithDefaultKeyStrict + , testProperty "findWithDefault is value-strict" pFindWithDefaultValueStrict + , testProperty "! is key-strict" $ keyStrict (flip (M.!)) + , testProperty "delete is key-strict" $ keyStrict M.delete + , testProperty "adjust is key-strict" pAdjustKeyStrict + , testProperty "adjust is value-strict" pAdjustValueStrict + , testProperty "insert is key-strict" pInsertKeyStrict + , testProperty "insert is value-strict" pInsertValueStrict + , testProperty "insertWith is key-strict" pInsertWithKeyStrict + , testProperty "insertWith is value-strict" pInsertWithValueStrict + , testProperty "insertLookupWithKey is key-strict" + pInsertLookupWithKeyKeyStrict + , testProperty "insertLookupWithKey is value-strict" + pInsertLookupWithKeyValueStrict + ] + ] + +------------------------------------------------------------------------ +-- * Test harness + +main :: IO () +main = defaultMain tests + +------------------------------------------------------------------------ +-- * Utilities + +keyStrict :: (Int -> Map Int Int -> a) -> Map Int Int -> Bool +keyStrict f m = isBottom $ f bottom m + +const2 :: a -> b -> c -> a +const2 x _ _ = x + +const3 :: a -> b -> c -> d -> a +const3 x _ _ _ = x diff --git a/tests/deprecated-properties.hs b/tests/deprecated-properties.hs new file mode 100644 --- /dev/null +++ b/tests/deprecated-properties.hs @@ -0,0 +1,102 @@ +{-# LANGUAGE CPP #-} + +-- This module tests the deprecated properties of Data.Map and Data.IntMap, +-- because these cannot be tested in either map-properties or +-- intmap-properties, as these modules are designed to work with the .Lazy and +-- .Strict modules. + +import qualified Data.Map as M +import qualified Data.Map.Strict as SM +import qualified Data.IntMap as IM +import qualified Data.IntMap.Strict as SIM + +import Test.Framework +import Test.Framework.Providers.QuickCheck2 +import Text.Show.Functions () + +default (Int) + +main :: IO () +main = defaultMain + [ testProperty "Data.Map.insertWith' as Strict.insertWith" prop_mapInsertWith'Strict + , testProperty "Data.Map.insertWith' undefined value" prop_mapInsertWith'Undefined + , testProperty "Data.Map.insertWithKey' as Strict.insertWithKey" prop_mapInsertWithKey'Strict + , testProperty "Data.Map.insertWithKey' undefined value" prop_mapInsertWithKey'Undefined + , testProperty "Data.Map.insertLookupWithKey' as Strict.insertLookupWithKey" prop_mapInsertLookupWithKey'Strict + , testProperty "Data.Map.insertLookupWithKey' undefined value" prop_mapInsertLookupWithKey'Undefined + , testProperty "Data.IntMap.insertWith' as Strict.insertWith" prop_intmapInsertWith'Strict + , testProperty "Data.IntMap.insertWith' undefined value" prop_intmapInsertWith'Undefined + , testProperty "Data.IntMap.insertWithKey' as Strict.insertWithKey" prop_intmapInsertWithKey'Strict + , testProperty "Data.IntMap.insertWithKey' undefined value" prop_intmapInsertWithKey'Undefined + ] + + +---------- Map properties ---------- + +prop_mapInsertWith'Strict :: [(Int, Int)] -> (Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_mapInsertWith'Strict xs f kxxs = + let m = M.fromList xs + insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + in insertList M.insertWith' == insertList SM.insertWith + +prop_mapInsertWith'Undefined :: [(Int, Int)] -> Bool +prop_mapInsertWith'Undefined xs = + let m = M.fromList xs + f _ x = x * 33 + insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs + in insertList M.insertWith' == insertList M.insertWith + +prop_mapInsertWithKey'Strict :: [(Int, Int)] -> (Int -> Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_mapInsertWithKey'Strict xs f kxxs = + let m = M.fromList xs + insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + in insertList M.insertWithKey' == insertList SM.insertWithKey + +prop_mapInsertWithKey'Undefined :: [(Int, Int)] -> Bool +prop_mapInsertWithKey'Undefined xs = + let m = M.fromList xs + f k _ x = (k + x) * 33 + insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs + in insertList M.insertWithKey' == insertList M.insertWithKey + +prop_mapInsertLookupWithKey'Strict :: [(Int, Int)] -> (Int -> Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_mapInsertLookupWithKey'Strict xs f kxxs = + let m = M.fromList xs + insertLookupList insLkp = scanr (\(kx, x) (_, mp) -> insLkp f kx x mp) (Nothing, m) kxxs + in insertLookupList M.insertLookupWithKey' == insertLookupList SM.insertLookupWithKey + +prop_mapInsertLookupWithKey'Undefined :: [(Int, Int)] -> Bool +prop_mapInsertLookupWithKey'Undefined xs = + let m = M.fromList xs + f k _ x = (k + x) * 33 + insertLookupList insLkp = scanr (\(kx, _) (_, mp) -> insLkp f kx undefined mp) (Nothing, m) xs + in insertLookupList M.insertLookupWithKey' == insertLookupList M.insertLookupWithKey + + +---------- IntMap properties ---------- + +prop_intmapInsertWith'Strict :: [(Int, Int)] -> (Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_intmapInsertWith'Strict xs f kxxs = + let m = IM.fromList xs + insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + in insertList IM.insertWith' == insertList SIM.insertWith + +prop_intmapInsertWith'Undefined :: [(Int, Int)] -> Bool +prop_intmapInsertWith'Undefined xs = + let m = IM.fromList xs + f _ x = x * 33 + insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs + in insertList IM.insertWith' == insertList IM.insertWith + +prop_intmapInsertWithKey'Strict :: [(Int, Int)] -> (Int -> Int -> Int -> Int) -> [(Int, Int)] -> Bool +prop_intmapInsertWithKey'Strict xs f kxxs = + let m = IM.fromList xs + insertList ins = foldr (\(kx, x) -> ins f kx x) m kxxs + in insertList IM.insertWithKey' == insertList SIM.insertWithKey + +prop_intmapInsertWithKey'Undefined :: [(Int, Int)] -> Bool +prop_intmapInsertWithKey'Undefined xs = + let m = IM.fromList xs + f k _ x = (k + x) * 33 + insertList ins = foldr (\(kx, _) -> ins f kx undefined) m xs + in insertList IM.insertWithKey' == insertList IM.insertWithKey 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,1056 @@ +{-# 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 = defaultMain + [ + 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 "fromList" prop_fromList + , 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 "splitRoot" prop_splitRoot + , testProperty "foldr" prop_foldr + , testProperty "foldr'" prop_foldr' + , testProperty "foldl" prop_foldl + , testProperty "foldl'" prop_foldl' + , testProperty "keysSet" prop_keysSet + , testProperty "fromSet" prop_fromSet + ] + +{-------------------------------------------------------------------- + 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_fromList :: [Int] -> Bool +prop_fromList xs + = case fromList (zip xs xs) of + t -> t == fromAscList (zip sort_xs sort_xs) && + t == fromDistinctAscList (zip nub_sort_xs nub_sort_xs) && + t == List.foldr (uncurry insert) empty (zip xs xs) + where sort_xs = sort xs + nub_sort_xs = List.map List.head $ List.group sort_xs + +---------------------------------------------------------------- + +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_splitRoot :: IMap -> Bool +prop_splitRoot s = loop ls && (s == unions ls) + where + ls = splitRoot s + loop [] = True + loop (s1:rst) = List.null + [ (x,y) | x <- toList s1 + , y <- toList (unions rst) + , x > y ] + +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,337 @@ +{-# LANGUAGE CPP #-} +#if MIN_VERSION_base(4,5,0) +import Data.Bits ((.&.), popCount) +import Data.Word (Word) +#else +import Data.Bits ((.&.)) +#endif +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 = defaultMain [ 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_splitRoot" prop_splitRoot + , testProperty "prop_partition" prop_partition + , testProperty "prop_filter" prop_filter +#if MIN_VERSION_base(4,5,0) + , testProperty "prop_bitcount" prop_bitcount +#endif + ] + +---------------------------------------------------------------- +-- 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_splitRoot :: IntSet -> Bool +prop_splitRoot s = loop ls && (s == unions ls) + where + ls = splitRoot s + loop [] = True + loop (s1:rst) = List.null + [ (x,y) | x <- toList s1 + , y <- toList (unions rst) + , x > y ] + +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) + +#if MIN_VERSION_base(4,5,0) +prop_bitcount :: Int -> Word -> Bool +prop_bitcount a w = bitcount_orig a w == bitcount_new a w + where + bitcount_orig a0 x0 = go a0 x0 + where go a 0 = a + go a x = go (a + 1) (x .&. (x-1)) + bitcount_new a x = a + popCount x +#endif 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,1203 @@ +{-# 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 = defaultMain + [ 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 "valid" prop_valid + , 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 "splitRoot" prop_splitRoot + , testProperty "split then link" prop_link + , 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 "fromList" prop_fromList + , 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 + ] + +{-------------------------------------------------------------------- + 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_valid :: UMap -> Bool +prop_valid 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_splitRoot :: UMap -> Bool +prop_splitRoot s = loop ls && (s == unions ls) + where + ls = splitRoot s + loop [] = True + loop (s1:rst) = List.null + [ (x,y) | x <- toList s1 + , y <- toList (unions rst) + , x > y ] + +prop_link :: Int -> UMap -> Bool +prop_link k t = let (l,r) = split k t + in valid (link k () l r) + +prop_merge :: Int -> UMap -> Bool +prop_merge k t = let (l,r) = split k t + in valid (merge l r) + +---------------------------------------------------------------- + +prop_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_fromList :: [Int] -> Bool +prop_fromList xs + = case fromList (zip xs xs) of + t -> t == fromAscList (zip sort_xs sort_xs) && + t == fromDistinctAscList (zip nub_sort_xs nub_sort_xs) && + t == List.foldr (uncurry insert) empty (zip xs xs) + where sort_xs = sort xs + nub_sort_xs = List.map List.head $ List.group sort_xs + +---------------------------------------------------------------- + +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,592 @@ +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 = defaultMain + [ 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 + ] + +------------------------------------------------------------------------ +-- 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,378 @@ +import qualified Data.IntSet as IntSet +import Data.List (nub,sort) +import qualified Data.List as List +import Data.Monoid (mempty) +import Data.Maybe +import Data.Set +import Prelude hiding (lookup, null, map, filter, foldr, foldl) +import Test.Framework +import Test.Framework.Providers.HUnit +import Test.Framework.Providers.QuickCheck2 +import Test.HUnit hiding (Test, Testable) +import Test.QuickCheck + +main :: IO () +main = defaultMain [ testCase "lookupLT" test_lookupLT + , testCase "lookupGT" test_lookupGT + , testCase "lookupLE" test_lookupLE + , testCase "lookupGE" test_lookupGE + , testCase "lookupIndex" test_lookupIndex + , testCase "findIndex" test_findIndex + , testCase "elemAt" test_elemAt + , testCase "deleteAt" test_deleteAt + , 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_Link" prop_Link + , 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_splitRoot" prop_splitRoot + , testProperty "prop_partition" prop_partition + , testProperty "prop_filter" prop_filter + ] + +---------------------------------------------------------------- +-- 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 + +{-------------------------------------------------------------------- + Indexed +--------------------------------------------------------------------} + +test_lookupIndex :: Assertion +test_lookupIndex = do + isJust (lookupIndex 2 (fromList [5,3])) @?= False + fromJust (lookupIndex 3 (fromList [5,3])) @?= 0 + fromJust (lookupIndex 5 (fromList [5,3])) @?= 1 + isJust (lookupIndex 6 (fromList [5,3])) @?= False + +test_findIndex :: Assertion +test_findIndex = do + findIndex 3 (fromList [5,3]) @?= 0 + findIndex 5 (fromList [5,3]) @?= 1 + +test_elemAt :: Assertion +test_elemAt = do + elemAt 0 (fromList [5,3]) @?= 3 + elemAt 1 (fromList [5,3]) @?= 5 + +test_deleteAt :: Assertion +test_deleteAt = do + deleteAt 0 (fromList [5,3]) @?= singleton 5 + deleteAt 1 (fromList [5,3]) @?= singleton 3 + +{-------------------------------------------------------------------- + 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_Link :: Int -> Property +prop_Link x = forValidUnitTree $ \t -> + let (l,r) = split x t + in valid (link 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_splitRoot :: Set Int -> Bool +prop_splitRoot s = loop ls && (s == unions ls) + where + ls = splitRoot s + loop [] = True + loop (s1:rst) = List.null + [ (x,y) | x <- toList s1 + , y <- toList (unions rst) + , x > y ] + +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)