diff --git a/changelog.md b/changelog.md --- a/changelog.md +++ b/changelog.md @@ -1,5 +1,207 @@ # Changelog for [`containers` package](http://github.com/haskell/containers) +## 0.8.1 *October 2026* + +### Additions + +* Add `compareSize` for `IntSet` and `IntMap`. (Soumik Sarkar) + ([#1135](https://github.com/haskell/containers/pull/1135), + [#1139](https://github.com/haskell/containers/pull/1139)) + +* Add `mapMaybe` for `Seq`, `Set` and `IntSet`. (Phil Hazelden) + ([#1159](https://github.com/haskell/containers/pull/1159)) + +* Add `fromSetA`, `fromSetMaybe`, and `fromSetMaybeA` for `Map` and `IntMap`. + (L0neGamer, Soumik Sarkar) + ([#1163](https://github.com/haskell/containers/pull/1163), + [#1165](https://github.com/haskell/containers/pull/1165), + [#1234](https://github.com/haskell/containers/pull/1234)) + +* Export `Tree` field selectors from `Data.Graph`. (Soumik Sarkar) + ([#1144](https://github.com/haskell/containers/pull/1144)) + +* Add `upsert`, `fromListUpsert`, `fromAscListUpsert`, and `fromDescListUpsert` + for `Map` and `IntMap`. (Soumik Sarkar) + ([#1145](https://github.com/haskell/containers/pull/1145), + [#1190](https://github.com/haskell/containers/pull/1190), + [#1199](https://github.com/haskell/containers/pull/1199)) + +* Add `pop` for `Map`, `Set`, `IntMap`, `IntSet`. (Soumik Sarkar) + ([#1152](https://github.com/haskell/containers/pull/1152)) + +* Add `Data.Sequence.toList`. (Soumik Sarkar) + ([#1192](https://github.com/haskell/containers/pull/1192)) + +* Add `fromDescList` for `IntSet` and `IntMap` (Soumik Sarkar) + ([#1194](https://github.com/haskell/containers/pull/1194)) + +* Add `takeR`, `dropR` and `splitAtR` for `Seq`. (Phil Crissman) + (see [#159](https://github.com/haskell/containers/issues/159)) + ([#1222](https://github.com/haskell/containers/pull/1222)) + +* Add `mapAssocsMonotonic` for `Map`. (Soumik Sarkar) + ([#1230](https://github.com/haskell/containers/pull/1230)) + +* Add `Data.Set.Merge`, a merge API for `Set`s. (Soumik Sarkar) + ([#1169](https://github.com/haskell/containers/pull/1169)) + +* Add `Data.Map.Merge.Set.Lazy` and `Data.Map.Merge.Set.Strict`, an API to + merge a `Map` and a `Set` into a `Map`. (Soumik Sarkar) + ([#1227](https://github.com/haskell/containers/pull/1227)) + +* Add the `dropMatched` and `whenMissing` merge strategies for `Map` and + `IntMap`. Add `filterA` for `Set`. (Soumik Sarkar) + ([#1240](https://github.com/haskell/containers/pull/1240), + [#1185](https://github.com/haskell/containers/pull/1185)) + +### Performance improvements + +* Improve performance of `Data.IntMap.fromAscList` and + `Data.IntSet.fromAscList`. (Soumik Sarkar) + ([#1123](https://github.com/haskell/containers/pull/1123)) + +* Improve performance of `Data.IntMap.fromList` and `Data.IntSet.fromList`. + (Soumik Sarkar) + ([#1129](https://github.com/haskell/containers/pull/1129), + [#1137](https://github.com/haskell/containers/pull/1137)) + +* Improved performance for `Data.IntMap.restrictKeys` and + `Data.IntMap.withoutKeys`. (Soumik Sarkar) + ([#1131](https://github.com/haskell/containers/pull/1131)) + +* Minor performance improvements for `IntMap` and `IntSet` by skipping some + unnecessary checks. (Soumik Sarkar) + ([#1136](https://github.com/haskell/containers/pull/1136)) + +* Improve performance of folds over `IntMap` and `IntSet`. (Soumik Sarkar) + ([#1149](https://github.com/haskell/containers/pull/1149)) + +* Improve performance of mapping for keys for `IntMap` and `IntSet`. + (Soumik Sarkar) + ([#1148](https://github.com/haskell/containers/pull/1148)) + +* Improve performance of `graphFromEdges`. (Soumik Sarkar) + ([#1151](https://github.com/haskell/containers/pull/1151)) + +* Improve performance of `Map`-`Map` and `Set`-`Set` operations. + (Soumik Sarkar) + ([#1141](https://github.com/haskell/containers/pull/1141)) + +* Improve performance of `Set` intersection. (Soumik Sarkar) + ([#1170](https://github.com/haskell/containers/pull/1170), + [#1172](https://github.com/haskell/containers/pull/1172)) + +* Improve performance of `nubOrdOn` and `nubIntOn`. (Soumik Sarkar) + ([#1206](https://github.com/haskell/containers/pull/1206), + [#1228](https://github.com/haskell/containers/pull/1228), + [#1229](https://github.com/haskell/containers/pull/1229)) + +* Use a different strategy for `Data.Set.alterF`, improving performance in + typical scenarios. (Soumik Sarkar) + ([#1215](https://github.com/haskell/containers/pull/1215)) + +* Reduce allocations when using `Data.Map.alterF`. (Soumik Sarkar) + ([#1219](https://github.com/haskell/containers/pull/1219)) + +* Improve performance of `Data.Tree`'s `leaves`, `edges`, and `foldr` for + `PostOrder`. (Soumik Sarkar) + ([#1245](https://github.com/haskell/containers/pull/1245)) + +* Allow specialization of `Data.Tree`'s `unfoldTreeM`, `unfoldForestM`, + `unfoldTreeM_BF`, and `unfoldForestM_BF`. (Soumik Sarkar) + ([#1257](https://github.com/haskell/containers/pull/1257)) + +* Improve performance of `Data.Tree.PostOrder`'s `foldl`, `foldr'`, `foldlMap1`, + `foldrMap1'`. (Soumik Sarkar) + ([#1259](https://github.com/haskell/containers/pull/1259)) + +### Documentation + +* Update contributing instructions. (Soumik Sarkar) + ([#1125](https://github.com/haskell/containers/pull/1125), + [#1150](https://github.com/haskell/containers/pull/1150), + [#1225](https://github.com/haskell/containers/pull/1225)) + +* Add and improve documentation (Jonathan Knowles, Soumik Sarkar, Tom Smeding, + Alexey Kuleshevich, RikuMinamiyama, Steve Shuck) + ([#1127](https://github.com/haskell/containers/pull/1127), + [#1138](https://github.com/haskell/containers/pull/1138), + [#1140](https://github.com/haskell/containers/pull/1140), + [#1143](https://github.com/haskell/containers/pull/1143), + [#1158](https://github.com/haskell/containers/pull/1158), + [#1164](https://github.com/haskell/containers/pull/1164), + [#1161](https://github.com/haskell/containers/pull/1161), + [#1168](https://github.com/haskell/containers/pull/1168), + [#1179](https://github.com/haskell/containers/pull/1179), + [#1189](https://github.com/haskell/containers/pull/1189), + [#1204](https://github.com/haskell/containers/pull/1204), + [#1218](https://github.com/haskell/containers/pull/1218), + [#1216](https://github.com/haskell/containers/pull/1216), + [#1231](https://github.com/haskell/containers/pull/1231), + [#1235](https://github.com/haskell/containers/pull/1235)) + +### Miscellaneous/internal + +* Fix bounds for `deepseq`. (Soumik Sarkar) + ([#1119](https://github.com/haskell/containers/pull/1119)) + +* Remove redundant `mappend` definitions in preparation for + [CLC #328](https://github.com/haskell/core-libraries-committee/issues/328). + (Soumik Sarkar) + ([#1142](https://github.com/haskell/containers/pull/1142)) + +* CI maintenance and improvements. (Soumik Sarkar, Lennart Augustsson) + ([#1147](https://github.com/haskell/containers/pull/1147), + [#1173](https://github.com/haskell/containers/pull/1173), + [#1177](https://github.com/haskell/containers/pull/1177), + [#1183](https://github.com/haskell/containers/pull/1183), + [#1180](https://github.com/haskell/containers/pull/1180), + [#1196](https://github.com/haskell/containers/pull/1196), + [#1207](https://github.com/haskell/containers/pull/1207), + [#1224](https://github.com/haskell/containers/pull/1224), + [#1236](https://github.com/haskell/containers/pull/1236), + [#1249](https://github.com/haskell/containers/pull/1249)) + +* Miscellaneous internal improvements. (Soumik Sarkar, Simon Hengel, konsumlamm) + ([#1126](https://github.com/haskell/containers/pull/1126), + [#1167](https://github.com/haskell/containers/pull/1167), + [#1175](https://github.com/haskell/containers/pull/1175), + [#1211](https://github.com/haskell/containers/pull/1211), + [#1210](https://github.com/haskell/containers/pull/1210), + [#1212](https://github.com/haskell/containers/pull/1212), + [#1213](https://github.com/haskell/containers/pull/1213), + [#1217](https://github.com/haskell/containers/pull/1217), + [#1223](https://github.com/haskell/containers/pull/1223), + [#1246](https://github.com/haskell/containers/pull/1246), + [#970](https://github.com/haskell/containers/pull/970)) + +* Additional exports from `Data.Set.Internal`. (Frank Staals) + ([#1178](https://github.com/haskell/containers/pull/1178)) + +* Test and benchmark improvements. (Soumik Sarkar, Alexandre Esteves) + ([#1181](https://github.com/haskell/containers/pull/1181), + [#1182](https://github.com/haskell/containers/pull/1182), + [#1188](https://github.com/haskell/containers/pull/1188), + [#1191](https://github.com/haskell/containers/pull/1191), + [#1198](https://github.com/haskell/containers/pull/1198), + [#1197](https://github.com/haskell/containers/pull/1197), + [#1203](https://github.com/haskell/containers/pull/1203), + [#1254](https://github.com/haskell/containers/pull/1254), + [#1255](https://github.com/haskell/containers/pull/1255)) + +* Use template-haskell-lift for GHC>=9.14 (Teo Camarasu) + ([#1162](https://github.com/haskell/containers/pull/1162)) + +* Drop redundant Applicative constraints. (Soumik Sarkar) + ([#1193](https://github.com/haskell/containers/pull/1193)) + +* Drop symlinks to make development easier on Windows. (AndreasPK) + ([#886](https://github.com/haskell/containers/pull/886)) + +* Expose unfoldings of some functions to make it possible to force-inline them. + (Soumik Sarkar) + ([#1252](https://github.com/haskell/containers/pull/1252)) + ## 0.8 *March 2025* ### Breaking changes diff --git a/containers.cabal b/containers.cabal --- a/containers.cabal +++ b/containers.cabal @@ -1,6 +1,6 @@ cabal-version: 2.2 name: containers -version: 0.8 +version: 0.8.1 license: BSD-3-Clause license-file: LICENSE maintainer: libraries@haskell.org @@ -30,7 +30,7 @@ tested-with: GHC ==8.2.2 || ==8.4.4 || ==8.6.5 || ==8.8.4 || ==8.10.7 || ==9.0.2 || ==9.2.8 || - ==9.4.8 || ==9.6.6 || ==9.8.4 || ==9.10.1 || ==9.12.1 + ==9.4.8 || ==9.6.7 || ==9.8.4 || ==9.10.3 || ==9.12.4 || ==9.14.1 source-repository head type: git @@ -38,8 +38,16 @@ Library default-language: Haskell2010 - build-depends: base >= 4.10 && < 5, array >= 0.4.0.0, deepseq >= 1.2 && < 1.6 - if impl(ghc) + build-depends: + base >= 4.10 && < 5 + , array >= 0.5.2.0 + , deepseq >= 1.4.3.0 && < 1.6 + -- template-haskell-lift was added as a boot library in GHC-9.14 + -- once we no longer wish to backport releases to older major releases, + -- this conditional can be dropped + if impl(ghc >= 9.14) + build-depends: template-haskell-lift >= 0.1 && <0.2 + elif impl(ghc) build-depends: template-haskell hs-source-dirs: src ghc-options: -O2 -Wall @@ -65,9 +73,13 @@ Data.Map.Strict.Internal Data.Map.Strict Data.Map.Merge.Strict + Data.Map.Merge.Set.Internal + Data.Map.Merge.Set.Lazy + Data.Map.Merge.Set.Strict Data.Map.Internal Data.Map.Internal.Debug Data.Set.Internal + Data.Set.Merge Data.Set Data.Graph Data.Sequence @@ -78,11 +90,10 @@ other-modules: Utils.Containers.Internal.Prelude Utils.Containers.Internal.State - Utils.Containers.Internal.StrictMaybe + Utils.Containers.Internal.Strict Utils.Containers.Internal.PtrEquality Utils.Containers.Internal.EqOrdUtil Utils.Containers.Internal.BitUtil Utils.Containers.Internal.BitQueue - Utils.Containers.Internal.StrictPair include-dirs: include diff --git a/src/Data/Containers/ListUtils.hs b/src/Data/Containers/ListUtils.hs --- a/src/Data/Containers/ListUtils.hs +++ b/src/Data/Containers/ListUtils.hs @@ -33,7 +33,7 @@ import qualified Data.IntSet as IntSet import Data.IntSet (IntSet) #ifdef __GLASGOW_HASKELL__ -import GHC.Exts ( build ) +import GHC.Exts (build, oneShot) #endif -- *** Ord-based nubbing *** @@ -72,9 +72,8 @@ -- -- @since 0.6.0.1 nubOrdOn :: Ord b => (a -> b) -> [a] -> [a] --- For some reason we need to write an explicit lambda here to allow this --- to inline when only applied to a function. -nubOrdOn f = \xs -> nubOrdOnExcluding f Set.empty xs +nubOrdOn f = -- Inline with 1 arg + \xs -> nubOrdOnExcluding f Set.empty xs {-# INLINE nubOrdOn #-} -- Splitting nubOrdOn like this means that we don't have to worry about @@ -82,12 +81,17 @@ nubOrdOnExcluding :: Ord b => (a -> b) -> Set b -> [a] -> [a] nubOrdOnExcluding f = go where - go _ [] = [] - go s (x:xs) - | fx `Set.member` s = go s xs - | otherwise = x : go (Set.insert fx s) xs + go !_ [] = [] + go !s (x:xs) = case tryInsertSet fx s of + Nothing -> go s xs + Just !s' -> -- See Note [Eager set insertions] + x : go s' xs where !fx = f x +tryInsertSet :: Ord a => a -> Set a -> Maybe (Set a) +tryInsertSet = Set.alterF (\found -> if found then Nothing else Just True) +{-# INLINE tryInsertSet #-} + #ifdef __GLASGOW_HASKELL__ -- We want this inlinable to specialize to the necessary Ord instance. {-# INLINABLE [1] nubOrdOnExcluding #-} @@ -110,14 +114,17 @@ -> (Set b -> r) -> Set b -> r -nubOrdOnFB f c x r s - | fx `Set.member` s = r s - | otherwise = x `c` r (Set.insert fx s) - where !fx = f x -{-# INLINABLE [0] nubOrdOnFB #-} +nubOrdOnFB f c = -- Inline with 2 args + \x r -> oneShot (\ !s -> + let !y = f x + in case tryInsertSet y s of + Nothing -> r s + Just !s' -> -- See Note [Eager set insertions] + x `c` r s') +{-# INLINE [0] nubOrdOnFB #-} constNubOn :: a -> b -> a -constNubOn x _ = x +constNubOn x !_ = x {-# INLINE [0] constNubOn #-} #endif @@ -153,9 +160,8 @@ -- -- @since 0.6.0.1 nubIntOn :: (a -> Int) -> [a] -> [a] --- For some reason we need to write an explicit lambda here to allow this --- to inline when only applied to a function. -nubIntOn f = \xs -> nubIntOnExcluding f IntSet.empty xs +nubIntOn f = -- Inline with 1 arg + \xs -> nubIntOnExcluding f IntSet.empty xs {-# INLINE nubIntOn #-} -- Splitting nubIntOn like this means that we don't have to worry about @@ -163,10 +169,12 @@ nubIntOnExcluding :: (a -> Int) -> IntSet -> [a] -> [a] nubIntOnExcluding f = go where - go _ [] = [] - go s (x:xs) + go !_ [] = [] + go !s (x:xs) | fx `IntSet.member` s = go s xs - | otherwise = x : go (IntSet.insert fx s) xs + | otherwise = + let !s' = IntSet.insert fx s -- See Note [Eager set insertions] + in x : go s' xs where !fx = f x #ifdef __GLASGOW_HASKELL__ @@ -189,9 +197,25 @@ -> (IntSet -> r) -> IntSet -> r -nubIntOnFB f c x r s - | fx `IntSet.member` s = r s - | otherwise = x `c` r (IntSet.insert fx s) - where !fx = f x -{-# INLINABLE [0] nubIntOnFB #-} +nubIntOnFB f c = -- Inline with 2 args + \x r -> oneShot (\ !s -> + let !y = f x + in if y `IntSet.member` s + then r s + else let !s' = IntSet.insert y s -- See Note [Eager set insertions] + in x `c` r s') +{-# INLINE [0] nubIntOnFB #-} #endif + +-- Note [Eager set insertions] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- +-- In nubOrd and nubInt we insert new elements into the set eagerly. This means +-- that we perform a bit of work before we yield the current element which is +-- not strictly necessary. +-- +-- The lazier option would be to create a thunk for the new set, which would get +-- forced by the membership check in the next step. However, a thunk has a small +-- overhead, and the small costs of thunks at every step adds up to a noticeable +-- amount of time and allocations overall. So, we avoid this and perform the +-- insertions eagerly instead. diff --git a/src/Data/Graph.hs b/src/Data/Graph.hs --- a/src/Data/Graph.hs +++ b/src/Data/Graph.hs @@ -6,9 +6,11 @@ {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveLift #-} {-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE Safe #-} {-# LANGUAGE TemplateHaskellQuotes #-} +#if !MIN_VERSION_array(0,5,7) +{-# LANGUAGE Trustworthy #-} #endif +#endif #ifdef DEFINE_PATTERN_SYNONYMS {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ViewPatterns #-} @@ -99,7 +101,8 @@ , flattenSCCs -- * Trees - , module Data.Tree + , Tree(..) + , Forest ) where @@ -107,17 +110,18 @@ import Prelude () #if USE_ST_MONAD import Control.Monad.ST -import Data.Array.ST.Safe (newArray, readArray, writeArray) +import Data.Array.ST (newArray, readArray, writeArray) # if USE_UNBOXED_ARRAYS -import Data.Array.ST.Safe (STUArray) +import Data.Array.ST (STUArray) # else -import Data.Array.ST.Safe (STArray) +import Data.Array.ST (STArray) # endif #else import Data.IntSet (IntSet) import qualified Data.IntSet as Set #endif -import Data.Tree (Tree(Node), Forest) +import Data.Tree (Tree(..), Forest) +import qualified Data.Tree as Tree -- std interfaces import Data.Foldable as F @@ -125,7 +129,6 @@ import qualified Data.Foldable1 as F1 #endif import Control.DeepSeq (NFData(rnf),NFData1(liftRnf)) -import Data.Maybe import Data.Array #if USE_UNBOXED_ARRAYS import qualified Data.Array.Unboxed as UA @@ -143,9 +146,13 @@ #ifdef __GLASGOW_HASKELL__ import GHC.Generics (Generic, Generic1) import Data.Data (Data) +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else import Language.Haskell.TH.Syntax (Lift(..)) -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () +# endif #endif -- Make sure we don't use Integer by mistake. @@ -191,7 +198,7 @@ deriving instance Generic (SCC vertex) -- There is no instance Lift (NonEmpty v) before template-haskell-2.15. -#if MIN_VERSION_template_haskell(2,15,0) +#if __GLASGOW_HASKELL__ > 808 -- | @since 0.6.6 deriving instance Lift vertex => Lift (SCC vertex) #else @@ -522,23 +529,30 @@ max_v = length edges0 - 1 bounds0 = (0,max_v) :: (Vertex, Vertex) sorted_edges = L.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 + graph = listArray bounds0 [keysToVertices ks | (_, _, ks) <- sorted_edges] + key_map = listArray bounds0 [k | (_, k, _) <- sorted_edges] + vertex_map = listArray bounds0 sorted_edges (_,k1,_) `lt` (_,k2,_) = k1 `compare` k2 - -- key_vertex :: key -> Maybe Vertex - -- returns Nothing for non-interesting vertices - key_vertex k = findVertex 0 max_v + keysToVertices = foldr f [] + where + f k vs = + let v = keyVertexGo k + in if v < 0 then vs else v:vs + + key_vertex k = + let v = keyVertexGo k + in if v < 0 then Nothing else Just v + + -- Binary search. Returns -1 when not found. + keyVertexGo k = findVertex 0 max_v where - findVertex a b | a > b - = Nothing + findVertex a b | a > b = -1 findVertex a b = case compare k (key_map ! mid) of LT -> findVertex a (mid-1) - EQ -> Just mid + EQ -> mid GT -> findVertex (mid+1) b where mid = a + (b - a) `div` 2 @@ -636,15 +650,6 @@ -- Algorithm 1: depth first search numbering ------------------------------------------------------------ -preorder' :: Tree a -> [a] -> [a] -preorder' (Node a ts) = (a :) . preorderF' ts - -preorderF' :: [Tree a] -> [a] -> [a] -preorderF' ts = foldr (.) id $ map preorder' ts - -preorderF :: [Tree a] -> [a] -preorderF ts = preorderF' ts [] - tabulate :: Bounds -> [Vertex] -> UArray Vertex Int tabulate bnds vs = UA.array bnds (zipWith (flip (,)) [1..] vs) -- Why zipWith (flip (,)) instead of just using zip with the @@ -653,21 +658,12 @@ -- list argument. preArr :: Bounds -> [Tree Vertex] -> UArray Vertex Int -preArr bnds = tabulate bnds . preorderF +preArr bnds = tabulate bnds . concatMap Tree.flatten ------------------------------------------------------------ -- Algorithm 2: topological sorting ------------------------------------------------------------ -postorder :: Tree a -> [a] -> [a] -postorder (Node a ts) = postorderF ts . (a :) - -postorderF :: [Tree a] -> [a] -> [a] -postorderF ts = foldr (.) id $ map postorder ts - -postOrd :: Graph -> [Vertex] -postOrd g = postorderF (dff g) [] - -- | \(O(V+E)\). 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. @@ -675,16 +671,22 @@ -- Note: A topological sort exists only when there are no cycles in the graph. -- If the graph has cycles, the output of this function will not be a -- topological sort. In such a case consider using 'scc'. -topSort :: Graph -> [Vertex] -topSort = reverse . postOrd +topSort :: Graph -> [Vertex] +topSort = reversePostOrder' . dff +-- Generates the result list at once. This is more efficient that being lazy if +-- we will consume the full result anyway. +reversePostOrder' :: [Tree a] -> [a] +reversePostOrder' = + F.foldl' (\xs t -> F.foldl' (flip (:)) xs (Tree.PostOrder t)) [] + -- | \(O(V+E)\). Reverse ordering of `topSort`. -- -- See note in 'topSort'. -- -- @since 0.6.4 reverseTopSort :: Graph -> [Vertex] -reverseTopSort = postOrd +reverseTopSort = concatMap (F.toList . Tree.PostOrder) . dff ------------------------------------------------------------ -- Algorithm 3: connected components @@ -710,8 +712,8 @@ -- > == [Node {rootLabel = 0, subForest = [Node {rootLabel = 1, subForest = [Node {rootLabel = 2, subForest = []}]}]} -- > ,Node {rootLabel = 3, subForest = []}] -scc :: Graph -> [Tree Vertex] -scc g = dfs g (reverse (postOrd (transposeG g))) +scc :: Graph -> [Tree Vertex] +scc g = dfs g (reversePostOrder' (dff (transposeG g))) ------------------------------------------------------------ -- Algorithm 5: Classifying edges @@ -753,7 +755,7 @@ -- -- > reachable (buildG (0,2) [(0,1), (1,2)]) 0 == [0,1,2] reachable :: Graph -> Vertex -> [Vertex] -reachable g v = preorderF (dfs g [v]) +reachable g v = concatMap Tree.flatten (dfs g [v]) -- | \(O(V+E)\). Returns @True@ if the second vertex reachable from the first. -- diff --git a/src/Data/IntMap.hs b/src/Data/IntMap.hs --- a/src/Data/IntMap.hs +++ b/src/Data/IntMap.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap diff --git a/src/Data/IntMap/Internal.hs b/src/Data/IntMap/Internal.hs --- a/src/Data/IntMap/Internal.hs +++ b/src/Data/IntMap/Internal.hs @@ -1,3886 +1,4549 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} -#ifdef __GLASGOW_HASKELL__ -{-# LANGUAGE DeriveLift #-} -{-# LANGUAGE MagicHash #-} -{-# LANGUAGE ScopedTypeVariables #-} -{-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE TypeFamilies #-} -{-# LANGUAGE Trustworthy #-} -#endif - -{-# OPTIONS_HADDOCK not-home #-} -{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-} - -#include "containers.h" - ------------------------------------------------------------------------------ --- | --- Module : Data.IntMap.Internal --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 --- (c) wren romano 2016 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Portability : portable --- --- = WARNING --- --- This module is considered __internal__. --- --- The Package Versioning Policy __does not apply__. --- --- The contents of this module may change __in any way whatsoever__ --- and __without any warning__ between minor versions of this package. --- --- Authors importing this module are expected to track development --- closely. --- --- --- = Finite Int Maps (lazy interface internals) --- --- The @'IntMap' v@ type represents a finite map (sometimes called a dictionary) --- from keys of type @Int@ to values of type @v@. --- --- --- == Implementation --- --- The implementation is based on /big-endian patricia trees/. This data --- structure performs especially well on binary operations like 'union' --- and 'intersection'. Additionally, 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, --- . --- --- @since 0.5.9 ------------------------------------------------------------------------------ - --- [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.Internal ( - -- * Map type - IntMap(..) -- instance Eq,Show - , Key - - -- * Operators - , (!), (!?), (\\) - - -- * Query - , null - , size - , member - , notMember - , lookup - , findWithDefault - , lookupLT - , lookupGT - , lookupLE - , lookupGE - , disjoint - - -- * Construction - , empty - , singleton - - -- ** Insertion - , insert - , insertWith - , insertWithKey - , insertLookupWithKey - - -- ** Delete\/Update - , delete - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - , alterF - - -- * Combine - - -- ** Union - , union - , unionWith - , unionWithKey - , unions - , unionsWith - - -- ** Difference - , difference - , differenceWith - , differenceWithKey - - -- ** Intersection - , intersection - , intersectionWith - , intersectionWithKey - - -- ** Symmetric difference - , symmetricDifference - - -- ** Compose - , compose - - -- ** General combining function - , SimpleWhenMissing - , SimpleWhenMatched - , runWhenMatched - , runWhenMissing - , merge - -- *** @WhenMatched@ tactics - , zipWithMaybeMatched - , zipWithMatched - -- *** @WhenMissing@ tactics - , mapMaybeMissing - , dropMissing - , preserveMissing - , mapMissing - , filterMissing - - -- ** Applicative general combining function - , WhenMissing (..) - , WhenMatched (..) - , mergeA - -- *** @WhenMatched@ tactics - -- | The tactics described for 'merge' work for - -- 'mergeA' as well. Furthermore, the following - -- are available. - , zipWithMaybeAMatched - , zipWithAMatched - -- *** @WhenMissing@ tactics - -- | The tactics described for 'merge' work for - -- 'mergeA' as well. Furthermore, the following - -- are available. - , traverseMaybeMissing - , traverseMissing - , filterAMissing - - -- ** Deprecated general combining function - , mergeWithKey - , mergeWithKey' - - -- * Traversal - -- ** Map - , map - , mapWithKey - , traverseWithKey - , traverseMaybeWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeys - , mapKeysWith - , mapKeysMonotonic - - -- * Folds - , foldr - , foldl - , foldrWithKey - , foldlWithKey - , foldMapWithKey - - -- ** Strict folds - , foldr' - , foldl' - , foldrWithKey' - , foldlWithKey' - - -- * Conversion - , elems - , keys - , assocs - , keysSet - , fromSet - - -- ** Lists - , toList - , fromList - , fromListWith - , fromListWithKey - - -- ** Ordered lists - , toAscList - , toDescList - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - - -- * Filter - , filter - , filterKeys - , filterWithKey - , restrictKeys - , withoutKeys - , partition - , partitionWithKey - - , takeWhileAntitone - , dropWhileAntitone - , spanAntitone - - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - - , split - , splitLookup - , splitRoot - - -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy - - -- * Min\/Max - , lookupMin - , lookupMax - , findMin - , findMax - , deleteMin - , deleteMax - , deleteFindMin - , deleteFindMax - , updateMin - , updateMax - , updateMinWithKey - , updateMaxWithKey - , minView - , maxView - , minViewWithKey - , maxViewWithKey - - -- * Debugging - , showTree - , showTreeWith - - -- * Utility - , link - , linkKey - , linkWithMask - , bin - , binCheckLeft - , binCheckRight - - -- * Used by "IntMap.Merge.Lazy" and "IntMap.Merge.Strict" - , mapWhenMissing - , mapWhenMatched - , lmapWhenMissing - , contramapFirstWhenMatched - , contramapSecondWhenMatched - , mapGentlyWhenMissing - , mapGentlyWhenMatched - ) where - -import Data.Functor.Identity (Identity (..)) -import Data.Semigroup (Semigroup(stimes)) -#if !(MIN_VERSION_base(4,11,0)) -import Data.Semigroup (Semigroup((<>))) -#endif -import Data.Semigroup (stimesIdempotentMonoid) -import Data.Functor.Classes - -import Control.DeepSeq (NFData(rnf),NFData1(liftRnf)) -import Data.Bits -import qualified Data.Foldable as Foldable -import Data.Maybe (fromMaybe) -import Utils.Containers.Internal.Prelude hiding - (lookup, map, filter, foldr, foldl, foldl', null) -import Prelude () - -import qualified Data.IntSet.Internal as IntSet -import Data.IntSet.Internal.IntTreeCommons - ( Key - , Prefix(..) - , nomatch - , left - , signBranch - , mask - , branchMask - , TreeTreeBranch(..) - , treeTreeBranch - , i2w - , Order(..) - ) -import Utils.Containers.Internal.BitUtil (shiftLL, shiftRL, iShiftRL) -import Utils.Containers.Internal.StrictPair - -#ifdef __GLASGOW_HASKELL__ -import Data.Coerce -import Data.Data (Data(..), Constr, mkConstr, constrIndex, - DataType, mkDataType, gcast1) -import qualified Data.Data as Data -import GHC.Exts (build) -import qualified GHC.Exts as GHCExts -import Language.Haskell.TH.Syntax (Lift) --- See Note [ Template Haskell Dependencies ] -import Language.Haskell.TH () -#endif -#if defined(__GLASGOW_HASKELL__) || defined(__MHS__) -import Text.Read -#endif -import qualified Control.Category as Category - - -{-------------------------------------------------------------------- - Types ---------------------------------------------------------------------} - - --- | A map of integers to values @a@. - --- See Note: Order of constructors -data IntMap a = Bin {-# UNPACK #-} !Prefix - !(IntMap a) - !(IntMap a) - | Tip {-# UNPACK #-} !Key a - | Nil - --- --- Note [IntMap structure and invariants] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- --- * Nil is never found as a child of Bin. --- --- * The Prefix of a Bin indicates the common high-order bits that all keys in --- the Bin share. --- --- * The least significant set bit of the Int value of a Prefix is called the --- mask bit. --- --- * All the bits to the left of the mask bit are called the shared prefix. All --- keys stored in the Bin begin with the shared prefix. --- --- * All keys in the left child of the Bin have the mask bit unset, and all keys --- in the right child have the mask bit set. It follows that --- --- 1. The Int value of the Prefix of a Bin is the smallest key that can be --- present in the right child of the Bin. --- --- 2. All keys in the right child of a Bin are greater than keys in the --- left child, with one exceptional situation. If the Bin separates --- negative and non-negative keys, the mask bit is the sign bit and the --- left child stores the non-negative keys while the right child stores the --- negative keys. --- --- * All bits to the right of the mask bit are set to 0 in a Prefix. --- - --- See Note [Okasaki-Gill] for how the implementation here relates to the one in --- Okasaki and Gill's paper. - --- Some stuff from "Data.IntSet.Internal", for 'restrictKeys' and --- 'withoutKeys' to use. -type IntSetPrefix = Int -type IntSetBitMap = Word - -#ifdef __GLASGOW_HASKELL__ --- | @since 0.6.6 -deriving instance Lift a => Lift (IntMap a) -#endif - -bitmapOf :: Int -> IntSetBitMap -bitmapOf x = shiftLL 1 (x .&. IntSet.suffixBitMask) -{-# INLINE bitmapOf #-} - -{-------------------------------------------------------------------- - 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 - --- | \(O(\min(n,W))\). Find the value at a key. --- Returns 'Nothing' when the element can not be found. --- --- > fromList [(5,'a'), (3,'b')] !? 1 == Nothing --- > fromList [(5,'a'), (3,'b')] !? 5 == Just 'a' --- --- @since 0.5.11 - -(!?) :: IntMap a -> Key -> Maybe a -(!?) m k = lookup 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 ---------------------------------------------------------------------} - --- | @mempty@ = 'empty' -instance Monoid (IntMap a) where - mempty = empty - mconcat = unions - mappend = (<>) - --- | @(<>)@ = 'union' --- --- @since 0.5.7 -instance Semigroup (IntMap a) where - (<>) = union - stimes = stimesIdempotentMonoid - --- | Folds in order of increasing key. -instance Foldable.Foldable IntMap where - fold = go - where go Nil = mempty - go (Tip _ v) = v - go (Bin p l r) - | signBranch p = go r `mappend` go l - | otherwise = 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 p l r) - | signBranch p = go r `mappend` go l - | otherwise = go l `mappend` go r - {-# INLINE foldMap #-} - foldl' = foldl' - {-# INLINE foldl' #-} - foldr' = foldr' - {-# INLINE foldr' #-} - length = size - {-# INLINE length #-} - null = null - {-# INLINE null #-} - toList = elems -- NB: Foldable.toList /= IntMap.toList - {-# INLINE toList #-} - elem = go - where go !_ Nil = False - go x (Tip _ y) = x == y - go x (Bin _ l r) = go x l || go x r - {-# INLINABLE elem #-} - maximum = start - where start Nil = error "Data.Foldable.maximum (for Data.IntMap): empty map" - start (Tip _ y) = y - start (Bin p l r) - | signBranch p = go (start r) l - | otherwise = go (start l) r - - go !m Nil = m - go m (Tip _ y) = max m y - go m (Bin _ l r) = go (go m l) r - {-# INLINABLE maximum #-} - minimum = start - where start Nil = error "Data.Foldable.minimum (for Data.IntMap): empty map" - start (Tip _ y) = y - start (Bin p l r) - | signBranch p = go (start r) l - | otherwise = go (start l) r - - go !m Nil = m - go m (Tip _ y) = min m y - go m (Bin _ l r) = go (go m l) r - {-# INLINABLE minimum #-} - sum = foldl' (+) 0 - {-# INLINABLE sum #-} - product = foldl' (*) 1 - {-# INLINABLE product #-} - --- | Traverses in order of increasing key. -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 - --- | @since 0.8 -instance NFData1 IntMap where - liftRnf rnfx = go - where - go Nil = () - go (Tip _ v) = rnfx v - go (Bin _ l r) = go l `seq` go 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" [] Data.Prefix - -intMapDataType :: DataType -intMapDataType = mkDataType "Data.IntMap.Internal.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 = go 0 - where - go !acc (Bin _ l r) = go (go acc l) r - go acc (Tip _ _) = 1 + acc - go acc Nil = acc - --- | \(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 = go - where - go (Bin p l r) - | nomatch k p = False - | left k p = go l - | otherwise = go r - go (Tip kx _) = k == kx - go Nil = False - --- | \(O(\min(n,W))\). Is the key not a member of the map? --- --- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False --- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True - -notMember :: Key -> IntMap a -> Bool -notMember k m = not $ member k m - --- | \(O(\min(n,W))\). Look up 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 = go - where - go (Bin p l r) | left k p = 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 = go - where - go (Bin p l r) | left k p = 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 = go - where - go (Bin p l r) | nomatch k p = def - | left k p = go l - | otherwise = go r - go (Tip kx x) | k == kx = x - | otherwise = def - go Nil = def - --- | \(O(\min(n,W))\). 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 = case t of - Bin p l r | signBranch p -> if k >= 0 then go r l else go Nil r - _ -> go Nil t - where - go def (Bin p l r) - | nomatch k p = if k < unPrefix p then unsafeFindMax def else unsafeFindMax r - | left k p = 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(\min(n,W))\). 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 = case t of - Bin p l r | signBranch p -> if k >= 0 then go Nil l else go l r - _ -> go Nil t - where - go def (Bin p l r) - | nomatch k p = if k < unPrefix p then unsafeFindMin l else unsafeFindMin def - | left k p = 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(\min(n,W))\). 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 = case t of - Bin p l r | signBranch p -> if k >= 0 then go r l else go Nil r - _ -> go Nil t - where - go def (Bin p l r) - | nomatch k p = if k < unPrefix p then unsafeFindMax def else unsafeFindMax r - | left k p = 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(\min(n,W))\). 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 = case t of - Bin p l r | signBranch p -> if k >= 0 then go Nil l else go l r - _ -> go Nil t - where - go def (Bin p l r) - | nomatch k p = if k < unPrefix p then unsafeFindMin l else unsafeFindMin def - | left k p = 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 - -{-------------------------------------------------------------------- - Disjoint ---------------------------------------------------------------------} --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- Check whether the key sets of two maps are disjoint --- (i.e. their 'intersection' is empty). --- --- > disjoint (fromList [(2,'a')]) (fromList [(1,()), (3,())]) == True --- > disjoint (fromList [(2,'a')]) (fromList [(1,'a'), (2,'b')]) == False --- > disjoint (fromList []) (fromList []) == True --- --- > disjoint a b == null (intersection a b) --- --- @since 0.6.2.1 -disjoint :: IntMap a -> IntMap b -> Bool -disjoint Nil _ = True -disjoint _ Nil = True -disjoint (Tip kx _) ys = notMember kx ys -disjoint xs (Tip ky _) = notMember ky xs -disjoint t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> disjoint l1 t2 - ABR -> disjoint r1 t2 - BAL -> disjoint t1 l2 - BAR -> disjoint t1 r2 - EQL -> disjoint l1 l2 && disjoint r1 r2 - NOM -> True - -{-------------------------------------------------------------------- - Compose ---------------------------------------------------------------------} --- | Relate the keys of one map to the values of --- the other, by using the values of the former as keys for lookups --- in the latter. --- --- Complexity: \( O(n * \min(m,W)) \), where \(m\) is the size of the first argument --- --- > compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) = fromList [(1,"A"),(2,"B")] --- --- @ --- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') --- @ --- --- __Note:__ Prior to v0.6.4, "Data.IntMap.Strict" exposed a version of --- 'compose' that forced the values of the output 'IntMap'. This version does --- not force these values. --- --- @since 0.6.3.1 -compose :: IntMap c -> IntMap Int -> IntMap c -compose bc !ab - | null bc = empty - | otherwise = mapMaybe (bc !?) ab - -{-------------------------------------------------------------------- - 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@(Bin p l r) - | nomatch k p = linkKey k (Tip k x) p t - | left k p = Bin p (insert k x l) r - | otherwise = Bin p l (insert k x r) -insert k x t@(Tip ky _) - | k==ky = Tip k x - | otherwise = link k (Tip k x) ky t -insert k x Nil = Tip k x - --- right-biased insertion, used by 'union' --- | \(O(\min(n,W))\). Insert with a combining function. --- @'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" --- --- Also see the performance note on 'fromListWith'. - -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" --- --- Also see the performance note on 'fromListWith'. - -insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWithKey f !k x t@(Bin p l r) - | nomatch k p = linkKey k (Tip k x) p t - | left k p = Bin p (insertWithKey f k x l) r - | otherwise = Bin p l (insertWithKey f k x r) -insertWithKey f k x t@(Tip ky y) - | k == ky = Tip k (f k x y) - | otherwise = link k (Tip k x) ky t -insertWithKey _ k x Nil = Tip k x - --- | \(O(\min(n,W))\). The expression (@'insertLookupWithKey' f k x map@) --- is a pair where the first element is equal to (@'lookup' k map@) --- 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")]) --- --- Also see the performance note on 'fromListWith'. - -insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) -insertLookupWithKey f !k x t@(Bin p l r) - | nomatch k p = (Nothing,linkKey k (Tip k x) p t) - | left k p = let (found,l') = insertLookupWithKey f k x l - in (found,Bin p l' r) - | otherwise = let (found,r') = insertLookupWithKey f k x r - in (found,Bin p l r') -insertLookupWithKey f k x t@(Tip ky y) - | k == ky = (Just y,Tip k (f k x y)) - | otherwise = (Nothing,link k (Tip k x) ky t) -insertLookupWithKey _ k x Nil = (Nothing,Tip k x) - - -{-------------------------------------------------------------------- - 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@(Bin p l r) - | nomatch k p = t - | left k p = binCheckLeft p (delete k l) r - | otherwise = binCheckRight p l (delete k r) -delete k t@(Tip ky _) - | k == ky = Nil - | otherwise = t -delete _k Nil = Nil - --- | \(O(\min(n,W))\). Adjust a value at a specific key. When the key is not --- a member of the map, the original map is returned. --- --- > 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 !k (Bin p l r) - | left k p = Bin p (adjustWithKey f k l) r - | otherwise = Bin p l (adjustWithKey f k r) -adjustWithKey f k t@(Tip ky y) - | k == ky = Tip ky (f k y) - | otherwise = t -adjustWithKey _ _ Nil = Nil - - --- | \(O(\min(n,W))\). The expression (@'update' f k map@) updates the value @x@ --- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is --- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. --- --- > 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 (Bin p l r) - | left k p = binCheckLeft p (updateWithKey f k l) r - | otherwise = binCheckRight p l (updateWithKey f k r) -updateWithKey f k t@(Tip ky y) - | k == ky = case (f k y) of - Just y' -> Tip ky y' - Nothing -> Nil - | otherwise = t -updateWithKey _ _ Nil = Nil - --- | \(O(\min(n,W))\). Look up and update. --- This function returns the 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 (Bin p l r) - | left k p = let !(found,l') = updateLookupWithKey f k l - in (found,binCheckLeft p l' r) - | otherwise = let !(found,r') = updateLookupWithKey f k r - in (found,binCheckRight p l r') -updateLookupWithKey f k t@(Tip ky y) - | k==ky = case (f k y) of - Just y' -> (Just y,Tip ky y') - Nothing -> (Just y,Nil) - | otherwise = (Nothing,t) -updateLookupWithKey _ _ Nil = (Nothing,Nil) - - - --- | \(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@(Bin p l r) - | nomatch k p = case f Nothing of - Nothing -> t - Just x -> linkKey k (Tip k x) p t - | left k p = binCheckLeft p (alter f k l) r - | otherwise = binCheckRight p l (alter f k r) -alter f k t@(Tip ky y) - | k==ky = case f (Just y) of - Just x -> Tip ky x - Nothing -> Nil - | otherwise = case f Nothing of - Just x -> link k (Tip k x) ky t - Nothing -> Tip ky y -alter f k Nil = case f Nothing of - Just x -> Tip k x - Nothing -> Nil - --- | \(O(\min(n,W))\). The expression (@'alterF' f k map@) alters the value @x@ at --- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, --- or update a value in an 'IntMap'. In short : @'lookup' k \<$\> 'alterF' f k m = f --- ('lookup' k m)@. --- --- Example: --- --- @ --- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) --- interactiveAlter k m = alterF f k m where --- f Nothing = do --- putStrLn $ show k ++ --- " was not found in the map. Would you like to add it?" --- getUserResponse1 :: IO (Maybe String) --- f (Just old) = do --- putStrLn $ "The key is currently bound to " ++ show old ++ --- ". Would you like to change or delete it?" --- getUserResponse2 :: IO (Maybe String) --- @ --- --- 'alterF' is the most general operation for working with an individual --- key that may or may not be in a given map. --- --- Note: 'alterF' is a flipped version of the @at@ combinator from --- @Control.Lens.At@. --- --- @since 0.5.8 - -alterF :: Functor f - => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) --- This implementation was stolen from 'Control.Lens.At'. -alterF f k m = (<$> f mv) $ \fres -> - case fres of - Nothing -> maybe m (const (delete k m)) mv - Just v' -> insert k v' m - where mv = lookup k m - -{-------------------------------------------------------------------- - Union ---------------------------------------------------------------------} --- | 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 :: Foldable f => f (IntMap a) -> IntMap a -unions xs - = Foldable.foldl' 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 :: Foldable f => (a->a->a) -> f (IntMap a) -> IntMap a -unionsWith f ts - = Foldable.foldl' (unionWith f) empty ts - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- The (left-biased) union of two maps. --- It prefers the first map when duplicate keys are encountered, --- i.e. (@'union' == 'unionWith' 'const'@). --- --- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] - -union :: IntMap a -> IntMap a -> IntMap a -union m1 m2 - = mergeWithKey' Bin const id id m1 m2 - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- The union with a combining function. --- --- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] --- --- Also see the performance note on 'fromListWith'. - -unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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")] --- --- Also see the performance note on 'fromListWith'. - -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(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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 - - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- Remove all the keys in a given set from a map. --- --- @ --- m \`withoutKeys\` s = 'filterWithKey' (\\k _ -> k ``IntSet.notMember`` s) m --- @ --- --- @since 0.5.8 -withoutKeys :: IntMap a -> IntSet.IntSet -> IntMap a -withoutKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> binCheckLeft p1 (withoutKeys l1 t2) r1 - ABR -> binCheckRight p1 l1 (withoutKeys r1 t2) - BAL -> withoutKeys t1 l2 - BAR -> withoutKeys t1 r2 - EQL -> bin p1 (withoutKeys l1 l2) (withoutKeys r1 r2) - NOM -> t1 - where -withoutKeys t1@(Bin p1 _ _) (IntSet.Tip p2 bm2) = - let px1 = unPrefix p1 - minbit = bitmapOf (px1 .&. (px1-1)) - lt_minbit = minbit - 1 - maxbit = bitmapOf (px1 .|. (px1-1)) - gt_maxbit = (-maxbit) `xor` maxbit - -- TODO(wrengr): should we manually inline/unroll 'updatePrefix' - -- and 'withoutBM' here, in order to avoid redundant case analyses? - in updatePrefix p2 t1 $ withoutBM (bm2 .|. lt_minbit .|. gt_maxbit) -withoutKeys t1@(Bin _ _ _) IntSet.Nil = t1 -withoutKeys t1@(Tip k1 _) t2 - | k1 `IntSet.member` t2 = Nil - | otherwise = t1 -withoutKeys Nil _ = Nil - - -updatePrefix - :: IntSetPrefix -> IntMap a -> (IntMap a -> IntMap a) -> IntMap a -updatePrefix !kp t@(Bin p l r) f - | unPrefix p .&. IntSet.suffixBitMask /= 0 = - if unPrefix p .&. IntSet.prefixBitMask == kp then f t else t - | nomatch kp p = t - | left kp p = binCheckLeft p (updatePrefix kp l f) r - | otherwise = binCheckRight p l (updatePrefix kp r f) -updatePrefix kp t@(Tip kx _) f - | kx .&. IntSet.prefixBitMask == kp = f t - | otherwise = t -updatePrefix _ Nil _ = Nil - - -withoutBM :: IntSetBitMap -> IntMap a -> IntMap a -withoutBM 0 t = t -withoutBM bm (Bin p l r) = - let leftBits = bitmapOf (unPrefix p) - 1 - bmL = bm .&. leftBits - bmR = bm `xor` bmL -- = (bm .&. complement leftBits) - in bin p (withoutBM bmL l) (withoutBM bmR r) -withoutBM bm t@(Tip k _) - -- TODO(wrengr): need we manually inline 'IntSet.Member' here? - | k `IntSet.member` IntSet.Tip (k .&. IntSet.prefixBitMask) bm = Nil - | otherwise = t -withoutBM _ Nil = Nil - - -{-------------------------------------------------------------------- - Intersection ---------------------------------------------------------------------} --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- The (left-biased) intersection of two maps (based on keys). --- --- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" - -intersection :: IntMap a -> IntMap b -> IntMap a -intersection m1 m2 - = mergeWithKey' bin const (const Nil) (const Nil) m1 m2 - - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- The restriction of a map to the keys in a set. --- --- @ --- m \`restrictKeys\` s = 'filterWithKey' (\\k _ -> k ``IntSet.member`` s) m --- @ --- --- @since 0.5.8 -restrictKeys :: IntMap a -> IntSet.IntSet -> IntMap a -restrictKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> restrictKeys l1 t2 - ABR -> restrictKeys r1 t2 - BAL -> restrictKeys t1 l2 - BAR -> restrictKeys t1 r2 - EQL -> bin p1 (restrictKeys l1 l2) (restrictKeys r1 r2) - NOM -> Nil -restrictKeys t1@(Bin p1 _ _) (IntSet.Tip p2 bm2) = - let px1 = unPrefix p1 - minbit = bitmapOf (px1 .&. (px1-1)) - ge_minbit = complement (minbit - 1) - maxbit = bitmapOf (px1 .|. (px1-1)) - le_maxbit = maxbit .|. (maxbit - 1) - -- TODO(wrengr): should we manually inline/unroll 'lookupPrefix' - -- and 'restrictBM' here, in order to avoid redundant case analyses? - in restrictBM (bm2 .&. ge_minbit .&. le_maxbit) (lookupPrefix p2 t1) -restrictKeys (Bin _ _ _) IntSet.Nil = Nil -restrictKeys t1@(Tip k1 _) t2 - | k1 `IntSet.member` t2 = t1 - | otherwise = Nil -restrictKeys Nil _ = Nil - - --- | \(O(\min(n,W))\). Restrict to the sub-map with all keys matching --- a key prefix. -lookupPrefix :: IntSetPrefix -> IntMap a -> IntMap a -lookupPrefix !kp t@(Bin p l r) - | unPrefix p .&. IntSet.suffixBitMask /= 0 = - if unPrefix p .&. IntSet.prefixBitMask == kp then t else Nil - | nomatch kp p = Nil - | left kp p = lookupPrefix kp l - | otherwise = lookupPrefix kp r -lookupPrefix kp t@(Tip kx _) - | (kx .&. IntSet.prefixBitMask) == kp = t - | otherwise = Nil -lookupPrefix _ Nil = Nil - - -restrictBM :: IntSetBitMap -> IntMap a -> IntMap a -restrictBM 0 _ = Nil -restrictBM bm (Bin p l r) = - let leftBits = bitmapOf (unPrefix p) - 1 - bmL = bm .&. leftBits - bmR = bm `xor` bmL -- = (bm .&. complement leftBits) - in bin p (restrictBM bmL l) (restrictBM bmR r) -restrictBM bm t@(Tip k _) - -- TODO(wrengr): need we manually inline 'IntSet.Member' here? - | k `IntSet.member` IntSet.Tip (k .&. IntSet.prefixBitMask) bm = t - | otherwise = Nil -restrictBM _ Nil = Nil - - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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 - -{-------------------------------------------------------------------- - Symmetric difference ---------------------------------------------------------------------} - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- The symmetric difference of two maps. --- --- The result contains entries whose keys appear in exactly one of the two maps. --- --- @ --- symmetricDifference --- (fromList [(0,\'q\'),(2,\'b\'),(4,\'w\'),(6,\'o\')]) --- (fromList [(0,\'e\'),(3,\'r\'),(6,\'t\'),(9,\'s\')]) --- == --- fromList [(2,\'b\'),(3,\'r\'),(4,\'w\'),(9,\'s\')] --- @ --- --- @since 0.8 -symmetricDifference :: IntMap a -> IntMap a -> IntMap a -symmetricDifference t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = - case treeTreeBranch p1 p2 of - ABL -> bin p1 (symmetricDifference l1 t2) r1 - ABR -> bin p1 l1 (symmetricDifference r1 t2) - BAL -> bin p2 (symmetricDifference t1 l2) r2 - BAR -> bin p2 l2 (symmetricDifference t1 r2) - EQL -> bin p1 (symmetricDifference l1 l2) (symmetricDifference r1 r2) - NOM -> link (unPrefix p1) t1 (unPrefix p2) t2 -symmetricDifference t1@(Bin _ _ _) t2@(Tip k2 _) = symDiffTip t2 k2 t1 -symmetricDifference t1@(Bin _ _ _) Nil = t1 -symmetricDifference t1@(Tip k1 _) t2 = symDiffTip t1 k1 t2 -symmetricDifference Nil t2 = t2 - -symDiffTip :: IntMap a -> Int -> IntMap a -> IntMap a -symDiffTip !t1 !k1 = go - where - go t2@(Bin p2 l2 r2) - | nomatch k1 p2 = linkKey k1 t1 p2 t2 - | left k1 p2 = bin p2 (go l2) r2 - | otherwise = bin p2 l2 (go r2) - go t2@(Tip k2 _) - | k1 == k2 = Nil - | otherwise = link k1 t1 k2 t2 - go Nil = t1 - -{-------------------------------------------------------------------- - MergeWithKey ---------------------------------------------------------------------} - --- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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'). --- --- __Warning__: 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@. - --- See Note [IntMap merge complexity] -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 -> 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 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> bin' p1 (go l1 t2) (g1 r1) - ABR -> bin' p1 (g1 l1) (go r1 t2) - BAL -> bin' p2 (go t1 l2) (g2 r2) - BAR -> bin' p2 (g2 l2) (go t1 r2) - EQL -> bin' p1 (go l1 l2) (go r1 r2) - NOM -> maybe_link (unPrefix p1) (g1 t1) (unPrefix p2) (g2 t2) - - go t1'@(Bin _ _ _) t2'@(Tip k2' _) = merge0 t2' k2' t1' - where - merge0 t2 k2 t1@(Bin p1 l1 r1) - | nomatch k2 p1 = maybe_link (unPrefix p1) (g1 t1) k2 (g2 t2) - | left k2 p1 = bin' p1 (merge0 t2 k2 l1) (g1 r1) - | otherwise = bin' p1 (g1 l1) (merge0 t2 k2 r1) - merge0 t2 k2 t1@(Tip k1 _) - | k1 == k2 = f t1 t2 - | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2) - merge0 t2 _ Nil = g2 t2 - - go t1@(Bin _ _ _) Nil = g1 t1 - - go t1'@(Tip k1' _) t2' = merge0 t1' k1' t2' - where - merge0 t1 k1 t2@(Bin p2 l2 r2) - | nomatch k1 p2 = maybe_link k1 (g1 t1) (unPrefix p2) (g2 t2) - | left k1 p2 = bin' p2 (merge0 t1 k1 l2) (g2 r2) - | otherwise = bin' p2 (g2 l2) (merge0 t1 k1 r2) - merge0 t1 k1 t2@(Tip k2 _) - | k1 == k2 = f t1 t2 - | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2) - merge0 t1 _ Nil = g1 t1 - - go Nil Nil = Nil - - go Nil t2 = g2 t2 - - maybe_link _ Nil _ t2 = t2 - maybe_link _ t1 _ Nil = t1 - maybe_link k1 t1 k2 t2 = link k1 t1 k2 t2 - {-# INLINE maybe_link #-} -{-# INLINE mergeWithKey' #-} - - -{-------------------------------------------------------------------- - mergeA ---------------------------------------------------------------------} - --- | A tactic for dealing with keys present in one map but not the --- other in 'merge' or 'mergeA'. --- --- A tactic of type @WhenMissing f k x z@ is an abstract representation --- of a function of type @Key -> x -> f (Maybe z)@. --- --- @since 0.5.9 - -data WhenMissing f x y = WhenMissing - { missingSubtree :: IntMap x -> f (IntMap y) - , missingKey :: Key -> x -> f (Maybe y)} - --- | @since 0.5.9 -instance (Applicative f, Monad f) => Functor (WhenMissing f x) where - fmap = mapWhenMissing - {-# INLINE fmap #-} - - --- | @since 0.5.9 -instance (Applicative f, Monad f) => Category.Category (WhenMissing f) - where - id = preserveMissing - f . g = - traverseMaybeMissing $ \ k x -> do - y <- missingKey g k x - case y of - Nothing -> pure Nothing - Just q -> missingKey f k q - {-# INLINE id #-} - {-# INLINE (.) #-} - - --- | Equivalent to @ReaderT k (ReaderT x (MaybeT f))@. --- --- @since 0.5.9 -instance (Applicative f, Monad f) => Applicative (WhenMissing f x) where - pure x = mapMissing (\ _ _ -> x) - f <*> g = - traverseMaybeMissing $ \k x -> do - res1 <- missingKey f k x - case res1 of - Nothing -> pure Nothing - Just r -> (pure $!) . fmap r =<< missingKey g k x - {-# INLINE pure #-} - {-# INLINE (<*>) #-} - - --- | Equivalent to @ReaderT k (ReaderT x (MaybeT f))@. --- --- @since 0.5.9 -instance (Applicative f, Monad f) => Monad (WhenMissing f x) where - m >>= f = - traverseMaybeMissing $ \k x -> do - res1 <- missingKey m k x - case res1 of - Nothing -> pure Nothing - Just r -> missingKey (f r) k x - {-# INLINE (>>=) #-} - - --- | Map covariantly over a @'WhenMissing' f x@. --- --- @since 0.5.9 -mapWhenMissing - :: (Applicative f, Monad f) - => (a -> b) - -> WhenMissing f x a - -> WhenMissing f x b -mapWhenMissing f t = WhenMissing - { missingSubtree = \m -> missingSubtree t m >>= \m' -> pure $! fmap f m' - , missingKey = \k x -> missingKey t k x >>= \q -> (pure $! fmap f q) } -{-# INLINE mapWhenMissing #-} - - --- | Map covariantly over a @'WhenMissing' f x@, using only a --- 'Functor f' constraint. -mapGentlyWhenMissing - :: Functor f - => (a -> b) - -> WhenMissing f x a - -> WhenMissing f x b -mapGentlyWhenMissing f t = WhenMissing - { missingSubtree = \m -> fmap f <$> missingSubtree t m - , missingKey = \k x -> fmap f <$> missingKey t k x } -{-# INLINE mapGentlyWhenMissing #-} - - --- | Map covariantly over a @'WhenMatched' f k x@, using only a --- 'Functor f' constraint. -mapGentlyWhenMatched - :: Functor f - => (a -> b) - -> WhenMatched f x y a - -> WhenMatched f x y b -mapGentlyWhenMatched f t = - zipWithMaybeAMatched $ \k x y -> fmap f <$> runWhenMatched t k x y -{-# INLINE mapGentlyWhenMatched #-} - - --- | Map contravariantly over a @'WhenMissing' f _ x@. --- --- @since 0.5.9 -lmapWhenMissing :: (b -> a) -> WhenMissing f a x -> WhenMissing f b x -lmapWhenMissing f t = WhenMissing - { missingSubtree = \m -> missingSubtree t (fmap f m) - , missingKey = \k x -> missingKey t k (f x) } -{-# INLINE lmapWhenMissing #-} - - --- | Map contravariantly over a @'WhenMatched' f _ y z@. --- --- @since 0.5.9 -contramapFirstWhenMatched - :: (b -> a) - -> WhenMatched f a y z - -> WhenMatched f b y z -contramapFirstWhenMatched f t = - WhenMatched $ \k x y -> runWhenMatched t k (f x) y -{-# INLINE contramapFirstWhenMatched #-} - - --- | Map contravariantly over a @'WhenMatched' f x _ z@. --- --- @since 0.5.9 -contramapSecondWhenMatched - :: (b -> a) - -> WhenMatched f x a z - -> WhenMatched f x b z -contramapSecondWhenMatched f t = - WhenMatched $ \k x y -> runWhenMatched t k x (f y) -{-# INLINE contramapSecondWhenMatched #-} - - --- | A tactic for dealing with keys present in one map but not the --- other in 'merge'. --- --- A tactic of type @SimpleWhenMissing x z@ is an abstract --- representation of a function of type @Key -> x -> Maybe z@. --- --- @since 0.5.9 -type SimpleWhenMissing = WhenMissing Identity - - --- | A tactic for dealing with keys present in both maps in 'merge' --- or 'mergeA'. --- --- A tactic of type @WhenMatched f x y z@ is an abstract representation --- of a function of type @Key -> x -> y -> f (Maybe z)@. --- --- @since 0.5.9 -newtype WhenMatched f x y z = WhenMatched - { matchedKey :: Key -> x -> y -> f (Maybe z) } - - --- | Along with zipWithMaybeAMatched, witnesses the isomorphism --- between @WhenMatched f x y z@ and @Key -> x -> y -> f (Maybe z)@. --- --- @since 0.5.9 -runWhenMatched :: WhenMatched f x y z -> Key -> x -> y -> f (Maybe z) -runWhenMatched = matchedKey -{-# INLINE runWhenMatched #-} - - --- | Along with traverseMaybeMissing, witnesses the isomorphism --- between @WhenMissing f x y@ and @Key -> x -> f (Maybe y)@. --- --- @since 0.5.9 -runWhenMissing :: WhenMissing f x y -> Key-> x -> f (Maybe y) -runWhenMissing = missingKey -{-# INLINE runWhenMissing #-} - - --- | @since 0.5.9 -instance Functor f => Functor (WhenMatched f x y) where - fmap = mapWhenMatched - {-# INLINE fmap #-} - - --- | @since 0.5.9 -instance (Monad f, Applicative f) => Category.Category (WhenMatched f x) - where - id = zipWithMatched (\_ _ y -> y) - f . g = - zipWithMaybeAMatched $ \k x y -> do - res <- runWhenMatched g k x y - case res of - Nothing -> pure Nothing - Just r -> runWhenMatched f k x r - {-# INLINE id #-} - {-# INLINE (.) #-} - - --- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ --- --- @since 0.5.9 -instance (Monad f, Applicative f) => Applicative (WhenMatched f x y) where - pure x = zipWithMatched (\_ _ _ -> x) - fs <*> xs = - zipWithMaybeAMatched $ \k x y -> do - res <- runWhenMatched fs k x y - case res of - Nothing -> pure Nothing - Just r -> (pure $!) . fmap r =<< runWhenMatched xs k x y - {-# INLINE pure #-} - {-# INLINE (<*>) #-} - - --- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ --- --- @since 0.5.9 -instance (Monad f, Applicative f) => Monad (WhenMatched f x y) where - m >>= f = - zipWithMaybeAMatched $ \k x y -> do - res <- runWhenMatched m k x y - case res of - Nothing -> pure Nothing - Just r -> runWhenMatched (f r) k x y - {-# INLINE (>>=) #-} - - --- | Map covariantly over a @'WhenMatched' f x y@. --- --- @since 0.5.9 -mapWhenMatched - :: Functor f - => (a -> b) - -> WhenMatched f x y a - -> WhenMatched f x y b -mapWhenMatched f (WhenMatched g) = - WhenMatched $ \k x y -> fmap (fmap f) (g k x y) -{-# INLINE mapWhenMatched #-} - - --- | A tactic for dealing with keys present in both maps in 'merge'. --- --- A tactic of type @SimpleWhenMatched x y z@ is an abstract --- representation of a function of type @Key -> x -> y -> Maybe z@. --- --- @since 0.5.9 -type SimpleWhenMatched = WhenMatched Identity - - --- | When a key is found in both maps, apply a function to the key --- and values and use the result in the merged map. --- --- > zipWithMatched --- > :: (Key -> x -> y -> z) --- > -> SimpleWhenMatched x y z --- --- @since 0.5.9 -zipWithMatched - :: Applicative f - => (Key -> x -> y -> z) - -> WhenMatched f x y z -zipWithMatched f = WhenMatched $ \ k x y -> pure . Just $ f k x y -{-# INLINE zipWithMatched #-} - - --- | When a key is found in both maps, apply a function to the key --- and values to produce an action and use its result in the merged --- map. --- --- @since 0.5.9 -zipWithAMatched - :: Applicative f - => (Key -> x -> y -> f z) - -> WhenMatched f x y z -zipWithAMatched f = WhenMatched $ \ k x y -> Just <$> f k x y -{-# INLINE zipWithAMatched #-} - - --- | When a key is found in both maps, apply a function to the key --- and values and maybe use the result in the merged map. --- --- > zipWithMaybeMatched --- > :: (Key -> x -> y -> Maybe z) --- > -> SimpleWhenMatched x y z --- --- @since 0.5.9 -zipWithMaybeMatched - :: Applicative f - => (Key -> x -> y -> Maybe z) - -> WhenMatched f x y z -zipWithMaybeMatched f = WhenMatched $ \ k x y -> pure $ f k x y -{-# INLINE zipWithMaybeMatched #-} - - --- | When a key is found in both maps, apply a function to the key --- and values, perform the resulting action, and maybe use the --- result in the merged map. --- --- This is the fundamental 'WhenMatched' tactic. --- --- @since 0.5.9 -zipWithMaybeAMatched - :: (Key -> x -> y -> f (Maybe z)) - -> WhenMatched f x y z -zipWithMaybeAMatched f = WhenMatched $ \ k x y -> f k x y -{-# INLINE zipWithMaybeAMatched #-} - - --- | Drop all the entries whose keys are missing from the other --- map. --- --- > dropMissing :: SimpleWhenMissing x y --- --- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing) --- --- but @dropMissing@ is much faster. --- --- @since 0.5.9 -dropMissing :: Applicative f => WhenMissing f x y -dropMissing = WhenMissing - { missingSubtree = const (pure Nil) - , missingKey = \_ _ -> pure Nothing } -{-# INLINE dropMissing #-} - - --- | Preserve, unchanged, the entries whose keys are missing from --- the other map. --- --- > preserveMissing :: SimpleWhenMissing x x --- --- prop> preserveMissing = Merge.Lazy.mapMaybeMissing (\_ x -> Just x) --- --- but @preserveMissing@ is much faster. --- --- @since 0.5.9 -preserveMissing :: Applicative f => WhenMissing f x x -preserveMissing = WhenMissing - { missingSubtree = pure - , missingKey = \_ v -> pure (Just v) } -{-# INLINE preserveMissing #-} - - --- | Map over the entries whose keys are missing from the other map. --- --- > mapMissing :: (k -> x -> y) -> SimpleWhenMissing x y --- --- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) --- --- but @mapMissing@ is somewhat faster. --- --- @since 0.5.9 -mapMissing :: Applicative f => (Key -> x -> y) -> WhenMissing f x y -mapMissing f = WhenMissing - { missingSubtree = \m -> pure $! mapWithKey f m - , missingKey = \k x -> pure $ Just (f k x) } -{-# INLINE mapMissing #-} - - --- | Map over the entries whose keys are missing from the other --- map, optionally removing some. This is the most powerful --- 'SimpleWhenMissing' tactic, but others are usually more efficient. --- --- > mapMaybeMissing :: (Key -> x -> Maybe y) -> SimpleWhenMissing x y --- --- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) --- --- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' --- operations. --- --- @since 0.5.9 -mapMaybeMissing - :: Applicative f => (Key -> x -> Maybe y) -> WhenMissing f x y -mapMaybeMissing f = WhenMissing - { missingSubtree = \m -> pure $! mapMaybeWithKey f m - , missingKey = \k x -> pure $! f k x } -{-# INLINE mapMaybeMissing #-} - - --- | Filter the entries whose keys are missing from the other map. --- --- > filterMissing :: (k -> x -> Bool) -> SimpleWhenMissing x x --- --- prop> filterMissing f = Merge.Lazy.mapMaybeMissing $ \k x -> guard (f k x) *> Just x --- --- but this should be a little faster. --- --- @since 0.5.9 -filterMissing - :: Applicative f => (Key -> x -> Bool) -> WhenMissing f x x -filterMissing f = WhenMissing - { missingSubtree = \m -> pure $! filterWithKey f m - , missingKey = \k x -> pure $! if f k x then Just x else Nothing } -{-# INLINE filterMissing #-} - - --- | Filter the entries whose keys are missing from the other map --- using some 'Applicative' action. --- --- > filterAMissing f = Merge.Lazy.traverseMaybeMissing $ --- > \k x -> (\b -> guard b *> Just x) <$> f k x --- --- but this should be a little faster. --- --- @since 0.5.9 -filterAMissing - :: Applicative f => (Key -> x -> f Bool) -> WhenMissing f x x -filterAMissing f = WhenMissing - { missingSubtree = \m -> filterWithKeyA f m - , missingKey = \k x -> bool Nothing (Just x) <$> f k x } -{-# INLINE filterAMissing #-} - - --- | \(O(n)\). Filter keys and values using an 'Applicative' predicate. -filterWithKeyA - :: Applicative f => (Key -> a -> f Bool) -> IntMap a -> f (IntMap a) -filterWithKeyA _ Nil = pure Nil -filterWithKeyA f t@(Tip k x) = (\b -> if b then t else Nil) <$> f k x -filterWithKeyA f (Bin p l r) - | signBranch p = liftA2 (flip (bin p)) (filterWithKeyA f r) (filterWithKeyA f l) - | otherwise = liftA2 (bin p) (filterWithKeyA f l) (filterWithKeyA f r) - --- | This wasn't in Data.Bool until 4.7.0, so we define it here -bool :: a -> a -> Bool -> a -bool f _ False = f -bool _ t True = t - - --- | Traverse over the entries whose keys are missing from the other --- map. --- --- @since 0.5.9 -traverseMissing - :: Applicative f => (Key -> x -> f y) -> WhenMissing f x y -traverseMissing f = WhenMissing - { missingSubtree = traverseWithKey f - , missingKey = \k x -> Just <$> f k x } -{-# INLINE traverseMissing #-} - - --- | Traverse over the entries whose keys are missing from the other --- map, optionally producing values to put in the result. This is --- the most powerful 'WhenMissing' tactic, but others are usually --- more efficient. --- --- @since 0.5.9 -traverseMaybeMissing - :: Applicative f => (Key -> x -> f (Maybe y)) -> WhenMissing f x y -traverseMaybeMissing f = WhenMissing - { missingSubtree = traverseMaybeWithKey f - , missingKey = f } -{-# INLINE traverseMaybeMissing #-} - - --- | \(O(n)\). Traverse keys\/values and collect the 'Just' results. --- --- @since 0.6.4 -traverseMaybeWithKey - :: Applicative f => (Key -> a -> f (Maybe b)) -> IntMap a -> f (IntMap b) -traverseMaybeWithKey f = go - where - go Nil = pure Nil - go (Tip k x) = maybe Nil (Tip k) <$> f k x - go (Bin p l r) - | signBranch p = liftA2 (flip (bin p)) (go r) (go l) - | otherwise = liftA2 (bin p) (go l) (go r) - - --- | Merge two maps. --- --- 'merge' takes two 'WhenMissing' tactics, a 'WhenMatched' tactic --- and two maps. It uses the tactics to merge the maps. Its behavior --- is best understood via its fundamental tactics, 'mapMaybeMissing' --- and 'zipWithMaybeMatched'. --- --- Consider --- --- @ --- merge (mapMaybeMissing g1) --- (mapMaybeMissing g2) --- (zipWithMaybeMatched f) --- m1 m2 --- @ --- --- Take, for example, --- --- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- 'merge' will first \"align\" these maps by key: --- --- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- It will then pass the individual entries and pairs of entries --- to @g1@, @g2@, or @f@ as appropriate: --- --- @ --- maybes = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"] --- @ --- --- This produces a 'Maybe' for each key: --- --- @ --- keys = 0 1 2 3 4 --- results = [Nothing, Just True, Just False, Nothing, Just True] --- @ --- --- Finally, the @Just@ results are collected into a map: --- --- @ --- return value = [(1, True), (2, False), (4, True)] --- @ --- --- The other tactics below are optimizations or simplifications of --- 'mapMaybeMissing' for special cases. Most importantly, --- --- * 'dropMissing' drops all the keys. --- * 'preserveMissing' leaves all the entries alone. --- --- When 'merge' is given three arguments, it is inlined at the call --- site. To prevent excessive inlining, you should typically use --- 'merge' to define your custom combining functions. --- --- --- Examples: --- --- prop> unionWithKey f = merge preserveMissing preserveMissing (zipWithMatched f) --- prop> intersectionWithKey f = merge dropMissing dropMissing (zipWithMatched f) --- prop> differenceWith f = merge diffPreserve diffDrop f --- prop> symmetricDifference = merge diffPreserve diffPreserve (\ _ _ _ -> Nothing) --- prop> mapEachPiece f g h = merge (diffMapWithKey f) (diffMapWithKey g) --- --- @since 0.5.9 -merge - :: SimpleWhenMissing a c -- ^ What to do with keys in @m1@ but not @m2@ - -> SimpleWhenMissing b c -- ^ What to do with keys in @m2@ but not @m1@ - -> SimpleWhenMatched a b c -- ^ What to do with keys in both @m1@ and @m2@ - -> IntMap a -- ^ Map @m1@ - -> IntMap b -- ^ Map @m2@ - -> IntMap c -merge g1 g2 f m1 m2 = - runIdentity $ mergeA g1 g2 f m1 m2 -{-# INLINE merge #-} - - --- | An applicative version of 'merge'. --- --- 'mergeA' takes two 'WhenMissing' tactics, a 'WhenMatched' --- tactic and two maps. It uses the tactics to merge the maps. --- Its behavior is best understood via its fundamental tactics, --- 'traverseMaybeMissing' and 'zipWithMaybeAMatched'. --- --- Consider --- --- @ --- mergeA (traverseMaybeMissing g1) --- (traverseMaybeMissing g2) --- (zipWithMaybeAMatched f) --- m1 m2 --- @ --- --- Take, for example, --- --- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3,\'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- 'mergeA' will first \"align\" these maps by key: --- --- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- It will then pass the individual entries and pairs of entries --- to @g1@, @g2@, or @f@ as appropriate: --- --- @ --- actions = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"] --- @ --- --- Next, it will perform the actions in the @actions@ list in order from --- left to right. --- --- @ --- keys = 0 1 2 3 4 --- results = [Nothing, Just True, Just False, Nothing, Just True] --- @ --- --- Finally, the @Just@ results are collected into a map: --- --- @ --- return value = [(1, True), (2, False), (4, True)] --- @ --- --- The other tactics below are optimizations or simplifications of --- 'traverseMaybeMissing' for special cases. Most importantly, --- --- * 'dropMissing' drops all the keys. --- * 'preserveMissing' leaves all the entries alone. --- * 'mapMaybeMissing' does not use the 'Applicative' context. --- --- When 'mergeA' is given three arguments, it is inlined at the call --- site. To prevent excessive inlining, you should generally only use --- 'mergeA' to define custom combining functions. --- --- @since 0.5.9 -mergeA - :: (Applicative f) - => WhenMissing f a c -- ^ What to do with keys in @m1@ but not @m2@ - -> WhenMissing f b c -- ^ What to do with keys in @m2@ but not @m1@ - -> WhenMatched f a b c -- ^ What to do with keys in both @m1@ and @m2@ - -> IntMap a -- ^ Map @m1@ - -> IntMap b -- ^ Map @m2@ - -> f (IntMap c) -mergeA - WhenMissing{missingSubtree = g1t, missingKey = g1k} - WhenMissing{missingSubtree = g2t, missingKey = g2k} - WhenMatched{matchedKey = f} - = go - where - go t1 Nil = g1t t1 - go Nil t2 = g2t t2 - - -- This case is already covered below. - -- go (Tip k1 x1) (Tip k2 x2) = mergeTips k1 x1 k2 x2 - - go (Tip k1 x1) t2' = merge2 t2' - where - merge2 t2@(Bin p2 l2 r2) - | nomatch k1 p2 = linkA k1 (subsingletonBy g1k k1 x1) (unPrefix p2) (g2t t2) - | left k1 p2 = binA p2 (merge2 l2) (g2t r2) - | otherwise = binA p2 (g2t l2) (merge2 r2) - merge2 (Tip k2 x2) = mergeTips k1 x1 k2 x2 - merge2 Nil = subsingletonBy g1k k1 x1 - - go t1' (Tip k2 x2) = merge1 t1' - where - merge1 t1@(Bin p1 l1 r1) - | nomatch k2 p1 = linkA (unPrefix p1) (g1t t1) k2 (subsingletonBy g2k k2 x2) - | left k2 p1 = binA p1 (merge1 l1) (g1t r1) - | otherwise = binA p1 (g1t l1) (merge1 r1) - merge1 (Tip k1 x1) = mergeTips k1 x1 k2 x2 - merge1 Nil = subsingletonBy g2k k2 x2 - - go t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> binA p1 (go l1 t2) (g1t r1) - ABR -> binA p1 (g1t l1) (go r1 t2) - BAL -> binA p2 (go t1 l2) (g2t r2) - BAR -> binA p2 (g2t l2) (go t1 r2) - EQL -> binA p1 (go l1 l2) (go r1 r2) - NOM -> linkA (unPrefix p1) (g1t t1) (unPrefix p2) (g2t t2) - - subsingletonBy :: Functor f => (Key -> a -> f (Maybe c)) -> Key -> a -> f (IntMap c) - subsingletonBy gk k x = maybe Nil (Tip k) <$> gk k x - {-# INLINE subsingletonBy #-} - - mergeTips k1 x1 k2 x2 - | k1 == k2 = maybe Nil (Tip k1) <$> f k1 x1 x2 - | k1 < k2 = liftA2 (subdoubleton k1 k2) (g1k k1 x1) (g2k k2 x2) - {- - = link_ k1 k2 <$> subsingletonBy g1k k1 x1 <*> subsingletonBy g2k k2 x2 - -} - | otherwise = liftA2 (subdoubleton k2 k1) (g2k k2 x2) (g1k k1 x1) - {-# INLINE mergeTips #-} - - subdoubleton _ _ Nothing Nothing = Nil - subdoubleton _ k2 Nothing (Just y2) = Tip k2 y2 - subdoubleton k1 _ (Just y1) Nothing = Tip k1 y1 - subdoubleton k1 k2 (Just y1) (Just y2) = link k1 (Tip k1 y1) k2 (Tip k2 y2) - {-# INLINE subdoubleton #-} - - -- | A variant of 'link_' which makes sure to execute side-effects - -- in the right order. - linkA - :: Applicative f - => Int -> f (IntMap a) - -> Int -> f (IntMap a) - -> f (IntMap a) - linkA k1 t1 k2 t2 - | i2w k1 < i2w k2 = binA p t1 t2 - | otherwise = binA p t2 t1 - where - m = branchMask k1 k2 - p = Prefix (mask k1 m .|. m) - {-# INLINE linkA #-} - - -- A variant of 'bin' that ensures that effects for negative keys are executed - -- first. - binA - :: Applicative f - => Prefix - -> f (IntMap a) - -> f (IntMap a) - -> f (IntMap a) - binA p a b - | signBranch p = liftA2 (flip (bin p)) b a - | otherwise = liftA2 (bin p) a b - {-# INLINE binA #-} -{-# INLINE mergeA #-} - - -{-------------------------------------------------------------------- - 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 l r | signBranch p -> binCheckRight p l (go f r) - _ -> go f t - where - go f' (Bin p l r) = binCheckLeft p (go f' l) r - go f' (Tip k y) = case f' k y of - Just y' -> Tip k y' - Nothing -> Nil - go _ Nil = 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 l r | signBranch p -> binCheckLeft p (go f l) r - _ -> go f t - where - go f' (Bin p l r) = binCheckRight p l (go f' r) - go f' (Tip k y) = case f' k y of - Just y' -> Tip k y' - Nothing -> Nil - go _ Nil = Nil - - -data View a = View {-# UNPACK #-} !Key a !(IntMap a) - --- | \(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 - _ -> Just $ case maxViewWithKeySure t of - View k v t' -> ((k, v), t') -{-# INLINE maxViewWithKey #-} - -maxViewWithKeySure :: IntMap a -> View a -maxViewWithKeySure t = - case t of - Nil -> error "maxViewWithKeySure Nil" - Bin p l r | signBranch p -> - case go l of View k a l' -> View k a (binCheckLeft p l' r) - _ -> go t - where - go (Bin p l r) = - case go r of View k a r' -> View k a (binCheckRight p l r') - go (Tip k y) = View k y Nil - go Nil = error "maxViewWithKey_go Nil" --- See note on NOINLINE at minViewWithKeySure -{-# NOINLINE maxViewWithKeySure #-} - --- | \(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 - _ -> Just $ case minViewWithKeySure t of - View k v t' -> ((k, v), t') --- We inline this to give GHC the best possible chance of --- getting rid of the Maybe, pair, and Int constructors, as --- well as a thunk under the Just. That is, we really want to --- be certain this inlines! -{-# INLINE minViewWithKey #-} - -minViewWithKeySure :: IntMap a -> View a -minViewWithKeySure t = - case t of - Nil -> error "minViewWithKeySure Nil" - Bin p l r | signBranch p -> - case go r of - View k a r' -> View k a (binCheckRight p l r') - _ -> go t - where - go (Bin p l r) = - case go l of View k a l' -> View k a (binCheckLeft p l' r) - go (Tip k y) = View k y Nil - go Nil = error "minViewWithKey_go Nil" --- There's never anything significant to be gained by inlining --- this. Sufficiently recent GHC versions will inline the wrapper --- anyway, which should be good enough. -{-# NOINLINE minViewWithKeySure #-} - --- | \(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) - --- | \(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 = fmap (\((_, x), t') -> (x, t')) (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 = fmap (\((_, x), t') -> (x, t')) (minViewWithKey t) - --- | \(O(\min(n,W))\). Delete and find the maximal element. --- This function throws an error if the map is empty. Use 'maxViewWithKey' --- if the map may be empty. -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. --- This function throws an error if the map is empty. Use 'minViewWithKey' --- if the map may be empty. -deleteFindMin :: IntMap a -> ((Key, a), IntMap a) -deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey - --- The KeyValue type is used when returning a key-value pair and helps with --- GHC optimizations. --- --- For lookupMinSure, if the return type is (Int, a), GHC compiles it to a --- worker $wlookupMinSure :: IntMap a -> (# Int, a #). If the return type is --- KeyValue a instead, the worker does not box the int and returns --- (# Int#, a #). --- For a modern enough GHC (>=9.4), this measure turns out to be unnecessary in --- this instance. We still use it for older GHCs and to make our intent clear. - -data KeyValue a = KeyValue {-# UNPACK #-} !Key a - -kvToTuple :: KeyValue a -> (Key, a) -kvToTuple (KeyValue k x) = (k, x) -{-# INLINE kvToTuple #-} - -lookupMinSure :: IntMap a -> KeyValue a -lookupMinSure (Tip k v) = KeyValue k v -lookupMinSure (Bin _ l _) = lookupMinSure l -lookupMinSure Nil = error "lookupMinSure Nil" - --- | \(O(\min(n,W))\). The minimal key of the map. Returns 'Nothing' if the map is empty. -lookupMin :: IntMap a -> Maybe (Key, a) -lookupMin Nil = Nothing -lookupMin (Tip k v) = Just (k,v) -lookupMin (Bin p l r) = - Just $! kvToTuple (lookupMinSure (if signBranch p then r else l)) -{-# INLINE lookupMin #-} -- See Note [Inline lookupMin] in Data.Set.Internal - --- | \(O(\min(n,W))\). The minimal key of the map. Calls 'error' if the map is empty. -findMin :: IntMap a -> (Key, a) -findMin t - | Just r <- lookupMin t = r - | otherwise = error "findMin: empty map has no minimal element" - -lookupMaxSure :: IntMap a -> KeyValue a -lookupMaxSure (Tip k v) = KeyValue k v -lookupMaxSure (Bin _ _ r) = lookupMaxSure r -lookupMaxSure Nil = error "lookupMaxSure Nil" - --- | \(O(\min(n,W))\). The maximal key of the map. Returns 'Nothing' if the map is empty. -lookupMax :: IntMap a -> Maybe (Key, a) -lookupMax Nil = Nothing -lookupMax (Tip k v) = Just (k,v) -lookupMax (Bin p l r) = - Just $! kvToTuple (lookupMaxSure (if signBranch p then l else r)) -{-# INLINE lookupMax #-} -- See Note [Inline lookupMin] in Data.Set.Internal - --- | \(O(\min(n,W))\). The maximal key of the map. Calls 'error' if the map is empty. -findMax :: IntMap a -> (Key, a) -findMax t - | Just r <- lookupMax t = r - | otherwise = error "findMax: empty map has no maximal element" - --- | \(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(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). - Is this a proper submap? (ie. a submap but not equal). - The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when - @keys m1@ and @keys 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 l1 r1) (Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> GT - ABR -> GT - BAL -> submapCmpLt l2 - BAR -> submapCmpLt r2 - EQL -> submapCmpEq - NOM -> GT -- disjoint - where - submapCmpLt t = case submapCmp predicate t1 t of - GT -> GT - _ -> LT - 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). --- Is this a submap? --- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). -isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool -isSubmapOf m1 m2 - = isSubmapOfBy (==) m1 m2 - -{- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). - 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 l1 r1) (Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> False - ABR -> False - BAL -> isSubmapOfBy predicate t1 l2 - BAR -> isSubmapOfBy predicate t1 r2 - EQL -> isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2 - NOM -> False -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 = go - where - go (Bin p l r) = Bin p (go l) (go r) - go (Tip k x) = Tip k (f x) - go Nil = Nil - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] map #-} -{-# RULES -"map/map" forall f g xs . map f (map g xs) = map (f . g) xs -"map/coerce" map coerce = coerce - #-} -#endif - --- | \(O(n)\). Map a function over all values in the map. --- --- > let f key x = (show key) ++ ":" ++ x --- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")] - -mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b -mapWithKey f t - = case t of - Bin p l r -> Bin p (mapWithKey f l) (mapWithKey f r) - Tip k x -> Tip k (f k x) - Nil -> Nil - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] mapWithKey #-} -{-# RULES -"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = - mapWithKey (\k a -> f k (g k a)) xs -"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = - mapWithKey (\k a -> f k (g a)) xs -"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = - mapWithKey (\k a -> f (g k a)) xs - #-} -#endif - --- | \(O(n)\). --- @'traverseWithKey' f 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 l r) - | signBranch p = liftA2 (flip (Bin p)) (go r) (go l) - | otherwise = liftA2 (Bin p) (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 l r - | signBranch p -> - let (a1,r') = mapAccumL f a r - (a2,l') = mapAccumL f a1 l - in (a2,Bin p l' r') - | otherwise -> - let (a1,l') = mapAccumL f a l - (a2,r') = mapAccumL f a1 r - in (a2,Bin p l' r') - Tip k x -> let (a',x') = f a k x in (a',Tip k x') - Nil -> (a,Nil) - --- | \(O(n)\). The function @'mapAccumRWithKey'@ 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 l r - | signBranch p -> - let (a1,l') = mapAccumRWithKey f a l - (a2,r') = mapAccumRWithKey f a1 r - in (a2,Bin p l' r') - | otherwise -> - let (a1,r') = mapAccumRWithKey f a r - (a2,l') = mapAccumRWithKey f a1 l - in (a2,Bin p 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" --- --- Also see the performance note on 'fromListWith'. - -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)\). --- @'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@. --- 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'. --- --- __Warning__: This function should be used only if @f@ is monotonically --- strictly increasing. This precondition is not checked. Use 'mapKeys' if the --- precondition may not hold. --- --- > 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 that satisfy some predicate. --- --- @ --- filterKeys p = 'filterWithKey' (\\k _ -> p k) --- @ --- --- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" --- --- @since 0.8 - -filterKeys :: (Key -> Bool) -> IntMap a -> IntMap a -filterKeys predicate = filterWithKey (\k _ -> predicate k) - --- | \(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 = go - where - go Nil = Nil - go t@(Tip k x) = if predicate k x then t else Nil - go (Bin p l r) = bin p (go l) (go r) - --- | \(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 l r -> - let (l1 :*: l2) = go predicate l - (r1 :*: r2) = go predicate r - in bin p l1 r1 :*: bin p l2 r2 - Tip k x - | predicate k x -> (t :*: Nil) - | otherwise -> (Nil :*: t) - Nil -> (Nil :*: Nil) - --- | \(O(\min(n,W))\). Take while a predicate on the keys holds. --- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@. --- See note at 'spanAntitone'. --- --- @ --- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' (p . fst) . 'toList' --- takeWhileAntitone p = 'filterWithKey' (\\k _ -> p k) --- @ --- --- @since 0.6.7 -takeWhileAntitone :: (Key -> Bool) -> IntMap a -> IntMap a -takeWhileAntitone predicate t = - case t of - Bin p l r - | signBranch p -> - if predicate 0 -- handle negative numbers. - then bin p (go predicate l) r - else go predicate r - _ -> go predicate t - where - go predicate' (Bin p l r) - | predicate' (unPrefix p) = bin p l (go predicate' r) - | otherwise = go predicate' l - go predicate' t'@(Tip ky _) - | predicate' ky = t' - | otherwise = Nil - go _ Nil = Nil - --- | \(O(\min(n,W))\). Drop while a predicate on the keys holds. --- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@. --- See note at 'spanAntitone'. --- --- @ --- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' (p . fst) . 'toList' --- dropWhileAntitone p = 'filterWithKey' (\\k _ -> not (p k)) --- @ --- --- @since 0.6.7 -dropWhileAntitone :: (Key -> Bool) -> IntMap a -> IntMap a -dropWhileAntitone predicate t = - case t of - Bin p l r - | signBranch p -> - if predicate 0 -- handle negative numbers. - then go predicate l - else bin p l (go predicate r) - _ -> go predicate t - where - go predicate' (Bin p l r) - | predicate' (unPrefix p) = go predicate' r - | otherwise = bin p (go predicate' l) r - go predicate' t'@(Tip ky _) - | predicate' ky = Nil - | otherwise = t' - go _ Nil = Nil - --- | \(O(\min(n,W))\). Divide a map at the point where a predicate on the keys stops holding. --- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@. --- --- @ --- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs) --- spanAntitone p xs = 'partitionWithKey' (\\k _ -> p k) xs --- @ --- --- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the map --- at some /unspecified/ point. --- --- @since 0.6.7 -spanAntitone :: (Key -> Bool) -> IntMap a -> (IntMap a, IntMap a) -spanAntitone predicate t = - case t of - Bin p l r - | signBranch p -> - if predicate 0 -- handle negative numbers. - then - case go predicate l of - (lt :*: gt) -> - let !lt' = bin p lt r - in (lt', gt) - else - case go predicate r of - (lt :*: gt) -> - let !gt' = bin p l gt - in (lt, gt') - _ -> case go predicate t of - (lt :*: gt) -> (lt, gt) - where - go predicate' (Bin p l r) - | predicate' (unPrefix p) - = case go predicate' r of (lt :*: gt) -> bin p l lt :*: gt - | otherwise - = case go predicate' l of (lt :*: gt) -> lt :*: bin p gt r - go predicate' t'@(Tip ky _) - | predicate' ky = (t' :*: Nil) - | otherwise = (Nil :*: t') - go _ 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 l r) - = bin p (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 l r) = - bin p l1 r1 :*: bin p 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 p l r - | signBranch p -> - if k >= 0 -- handle negative numbers. - then - case go k l of - (lt :*: gt) -> - let !lt' = bin p lt r - in (lt', gt) - else - case go k r of - (lt :*: gt) -> - let !gt' = bin p l gt - in (lt, gt') - _ -> case go k t of - (lt :*: gt) -> (lt, gt) - where - go !k' t'@(Bin p l r) - | nomatch k' p = if k' < unPrefix p then Nil :*: t' else t' :*: Nil - | left k' p = case go k' l of (lt :*: gt) -> lt :*: bin p gt r - | otherwise = case go k' r of (lt :*: gt) -> bin p l lt :*: gt - go k' t'@(Tip ky _) - | k' > ky = (t' :*: Nil) - | k' < ky = (Nil :*: t') - | otherwise = (Nil :*: Nil) - go _ Nil = (Nil :*: Nil) - - -data SplitLookup a = SplitLookup !(IntMap a) !(Maybe a) !(IntMap a) - -mapLT :: (IntMap a -> IntMap a) -> SplitLookup a -> SplitLookup a -mapLT f (SplitLookup lt fnd gt) = SplitLookup (f lt) fnd gt -{-# INLINE mapLT #-} - -mapGT :: (IntMap a -> IntMap a) -> SplitLookup a -> SplitLookup a -mapGT f (SplitLookup lt fnd gt) = SplitLookup lt fnd (f gt) -{-# INLINE mapGT #-} - --- | \(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 - case t of - Bin p l r - | signBranch p -> - if k >= 0 -- handle negative numbers. - then mapLT (flip (bin p) r) (go k l) - else mapGT (bin p l) (go k r) - _ -> go k t - of SplitLookup lt fnd gt -> (lt, fnd, gt) - where - go !k' t'@(Bin p l r) - | nomatch k' p = - if k' < unPrefix p - then SplitLookup Nil Nothing t' - else SplitLookup t' Nothing Nil - | left k' p = mapGT (flip (bin p) r) (go k' l) - | otherwise = mapLT (bin p l) (go k' r) - go k' t'@(Tip ky y) - | k' > ky = SplitLookup t' Nothing Nil - | k' < ky = SplitLookup Nil Nothing t' - | otherwise = SplitLookup Nil (Just y) Nil - go _ Nil = SplitLookup 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 p l r - | signBranch p -> 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 p l r - | signBranch p -> 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)\). 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 p l r - | signBranch p -> 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 p l r - | signBranch p -> 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)\). 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 p l r - | signBranch p -> 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 p l r - | signBranch p -> 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)\). 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 p l r - | signBranch p -> 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 p l r - | signBranch p -> 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)\). 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. --- --- @since 0.5.4 -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 p l r) - | signBranch p = go r `mappend` go l - | otherwise = 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)\). 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 l r) - | unPrefix p .&. IntSet.suffixBitMask == 0 - = IntSet.Bin p (keysSet l) (keysSet r) - | otherwise - = IntSet.Tip (unPrefix p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r) - where computeBm !acc (Bin _ l' r') = computeBm (computeBm acc l') r' - computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx - computeBm _ Nil = error "Data.IntSet.keysSet: Nil" - --- | \(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 l r) = Bin p (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 = case bits of - 0 -> Tip prefix (g prefix) - _ -> case bits `iShiftRL` 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 (prefix .|. bits2)) - (buildTree g prefix bmask bits2) - (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2) - -{-------------------------------------------------------------------- - Lists ---------------------------------------------------------------------} - -#ifdef __GLASGOW_HASKELL__ --- | @since 0.5.6.2 -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 - = Foldable.foldl' ins empty xs - where - ins t (k,x) = insert k x t - --- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. --- --- > fromListWith (++) [(5,"a"), (5,"b"), (3,"x"), (5,"c")] == fromList [(3, "x"), (5, "cba")] --- > fromListWith (++) [] == empty --- --- Note the reverse ordering of @"cba"@ in the example. --- --- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@. --- --- === Performance --- --- You should ensure that the given @f@ is fast with this order of arguments. --- --- Symmetric functions may be slow in one order, and fast in another. --- For the common case of collecting values of matching keys in a list, as above: --- --- The complexity of @(++) a b@ is \(O(a)\), so it is fast when given a short list as its first argument. --- Thus: --- --- > fromListWith (++) (replicate 1000000 (3, "x")) -- O(n), fast --- > fromListWith (flip (++)) (replicate 1000000 (3, "x")) -- O(n²), extremely slow --- --- because they evaluate as, respectively: --- --- > fromList [(3, "x" ++ ("x" ++ "xxxxx..xxxxx"))] -- O(n) --- > fromList [(3, ("xxxxx..xxxxx" ++ "x") ++ "x")] -- O(n²) --- --- Thus, to get good performance with an operation like @(++)@ while also preserving --- the same order as in the input list, reverse the input: --- --- > fromListWith (++) (reverse [(5,"a"), (5,"b"), (5,"c")]) == fromList [(5, "abc")] --- --- and it is always fast to combine singleton-list values @[v]@ with @fromListWith (++)@, as in: --- --- > fromListWith (++) $ reverse $ map (\(k, v) -> (k, [v])) someListOfTuples - -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 --- --- Also see the performance note on 'fromListWith'. - -fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWithKey f xs - = Foldable.foldl' 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. --- --- __Warning__: This function should be used only if the keys are in --- non-decreasing order. This precondition is not checked. Use 'fromList' if the --- precondition may not hold. --- --- > 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 = fromMonoListWithKey Nondistinct (\_ x _ -> x) -{-# NOINLINE fromAscList #-} - --- | \(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. --- --- __Warning__: This function should be used only if the keys are in --- non-decreasing order. This precondition is not checked. Use 'fromListWith' if --- the precondition may not hold. --- --- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] --- --- Also see the performance note on 'fromListWith'. - -fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f = fromMonoListWithKey Nondistinct (\_ x y -> f x y) -{-# NOINLINE fromAscListWith #-} - --- | \(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. --- --- __Warning__: This function should be used only if the keys are in --- non-decreasing order. This precondition is not checked. Use 'fromListWithKey' --- if the precondition may not hold. --- --- > 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")] --- --- Also see the performance note on 'fromListWith'. - -fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey f = fromMonoListWithKey Nondistinct f -{-# NOINLINE fromAscListWithKey #-} - --- | \(O(n)\). Build a map from a list of key\/value pairs where --- the keys are in ascending order and all distinct. --- --- __Warning__: This function should be used only if the keys are in --- strictly increasing order. This precondition is not checked. Use 'fromList' --- if the precondition may not hold. --- --- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] - -fromDistinctAscList :: [(Key,a)] -> IntMap a -fromDistinctAscList = fromMonoListWithKey Distinct (\_ x _ -> x) -{-# NOINLINE fromDistinctAscList #-} - --- | \(O(n)\). Build a map from a list of key\/value pairs with monotonic keys --- and a combining function. --- --- The precise conditions under which this function works are subtle: --- For any branch mask, keys with the same prefix w.r.t. the branch --- mask must occur consecutively in the list. --- --- Also see the performance note on 'fromListWith'. - -fromMonoListWithKey :: Distinct -> (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromMonoListWithKey distinct f = go - where - go [] = Nil - go ((kx,vx) : zs1) = addAll' kx vx zs1 - - -- `addAll'` collects all keys equal to `kx` into a single value, - -- and then proceeds with `addAll`. - addAll' !kx vx [] - = Tip kx vx - addAll' !kx vx ((ky,vy) : zs) - | Nondistinct <- distinct, kx == ky - = let v = f kx vy vx in addAll' ky v zs - -- inlined: | otherwise = addAll kx (Tip kx vx) (ky : zs) - | m <- branchMask kx ky - , Inserted ty zs' <- addMany' m ky vy zs - = addAll kx (linkWithMask m ky ty kx (Tip kx vx)) zs' - - -- for `addAll` and `addMany`, kx is /a/ key inside the tree `tx` - -- `addAll` consumes the rest of the list, adding to the tree `tx` - addAll !_kx !tx [] - = tx - addAll !kx !tx ((ky,vy) : zs) - | m <- branchMask kx ky - , Inserted ty zs' <- addMany' m ky vy zs - = addAll kx (linkWithMask m ky ty kx tx) zs' - - -- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`. - addMany' !_m !kx vx [] - = Inserted (Tip kx vx) [] - addMany' !m !kx vx zs0@((ky,vy) : zs) - | Nondistinct <- distinct, kx == ky - = let v = f kx vy vx in addMany' m ky v zs - -- inlined: | otherwise = addMany m kx (Tip kx vx) (ky : zs) - | mask kx m /= mask ky m - = Inserted (Tip kx vx) zs0 - | mxy <- branchMask kx ky - , Inserted ty zs' <- addMany' mxy ky vy zs - = addMany m kx (linkWithMask mxy ky ty kx (Tip kx vx)) zs' - - -- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `kx`. - addMany !_m !_kx tx [] - = Inserted tx [] - addMany !m !kx tx zs0@((ky,vy) : zs) - | mask kx m /= mask ky m - = Inserted tx zs0 - | mxy <- branchMask kx ky - , Inserted ty zs' <- addMany' mxy ky vy zs - = addMany m kx (linkWithMask mxy ky ty kx tx) zs' -{-# INLINE fromMonoListWithKey #-} - -data Inserted a = Inserted !(IntMap a) ![(Key,a)] - -data Distinct = Distinct | Nondistinct - -{-------------------------------------------------------------------- - Eq ---------------------------------------------------------------------} -instance Eq a => Eq (IntMap a) where - (==) = equal - -equal :: Eq a => IntMap a -> IntMap a -> Bool -equal (Bin p1 l1 r1) (Bin p2 l2 r2) - = (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 -{-# INLINABLE equal #-} - --- | @since 0.5.9 -instance Eq1 IntMap where - liftEq eq = go - where - go (Bin p1 l1 r1) (Bin p2 l2 r2) = p1 == p2 && go l1 l2 && go r1 r2 - go (Tip kx x) (Tip ky y) = kx == ky && eq x y - go Nil Nil = True - go _ _ = False - {-# INLINE liftEq #-} - -{-------------------------------------------------------------------- - Ord ---------------------------------------------------------------------} - -instance Ord a => Ord (IntMap a) where - compare m1 m2 = liftCmp compare m1 m2 - {-# INLINABLE compare #-} - --- | @since 0.5.9 -instance Ord1 IntMap where - liftCompare = liftCmp - -liftCmp :: (a -> b -> Ordering) -> IntMap a -> IntMap b -> Ordering -liftCmp cmp m1 m2 = case (splitSign m1, splitSign m2) of - ((l1, r1), (l2, r2)) -> case go l1 l2 of - A_LT_B -> LT - A_Prefix_B -> if null r1 then LT else GT - A_EQ_B -> case go r1 r2 of - A_LT_B -> LT - A_Prefix_B -> LT - A_EQ_B -> EQ - B_Prefix_A -> GT - A_GT_B -> GT - B_Prefix_A -> if null r2 then GT else LT - A_GT_B -> GT - where - go t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> case go l1 t2 of - A_Prefix_B -> A_GT_B - A_EQ_B -> B_Prefix_A - o -> o - ABR -> A_LT_B - BAL -> case go t1 l2 of - A_EQ_B -> A_Prefix_B - B_Prefix_A -> A_LT_B - o -> o - BAR -> A_GT_B - EQL -> case go l1 l2 of - A_Prefix_B -> A_GT_B - A_EQ_B -> go r1 r2 - B_Prefix_A -> A_LT_B - o -> o - NOM -> if unPrefix p1 < unPrefix p2 then A_LT_B else A_GT_B - go (Bin _ l1 _) (Tip k2 x2) = case lookupMinSure l1 of - KeyValue k1 x1 -> case compare k1 k2 <> cmp x1 x2 of - LT -> A_LT_B - EQ -> B_Prefix_A - GT -> A_GT_B - go (Tip k1 x1) (Bin _ l2 _) = case lookupMinSure l2 of - KeyValue k2 x2 -> case compare k1 k2 <> cmp x1 x2 of - LT -> A_LT_B - EQ -> A_Prefix_B - GT -> A_GT_B - go (Tip k1 x1) (Tip k2 x2) = case compare k1 k2 <> cmp x1 x2 of - LT -> A_LT_B - EQ -> A_EQ_B - GT -> A_GT_B - go Nil Nil = A_EQ_B - go Nil _ = A_Prefix_B - go _ Nil = B_Prefix_A -{-# INLINE liftCmp #-} - --- Split into negative and non-negative -splitSign :: IntMap a -> (IntMap a, IntMap a) -splitSign t@(Bin p l r) - | signBranch p = (r, l) - | unPrefix p < 0 = (t, Nil) - | otherwise = (Nil, t) -splitSign t@(Tip k _) - | k < 0 = (t, Nil) - | otherwise = (Nil, t) -splitSign Nil = (Nil, Nil) -{-# INLINE splitSign #-} - -{-------------------------------------------------------------------- - Functor ---------------------------------------------------------------------} - -instance Functor IntMap where - fmap = map - -#ifdef __GLASGOW_HASKELL__ - a <$ Bin p l r = Bin p (a <$ l) (a <$ r) - a <$ Tip k _ = Tip k a - _ <$ Nil = Nil -#endif - -{-------------------------------------------------------------------- - Show ---------------------------------------------------------------------} - -instance Show a => Show (IntMap a) where - showsPrec d m = showParen (d > 10) $ - showString "fromList " . shows (toList m) - --- | @since 0.5.9 -instance Show1 IntMap where - liftShowsPrec sp sl d m = - showsUnaryWith (liftShowsPrec sp' sl') "fromList" d (toList m) - where - sp' = liftShowsPrec sp sl - sl' = liftShowList sp sl - -{-------------------------------------------------------------------- - Read ---------------------------------------------------------------------} -instance (Read e) => Read (IntMap e) where -#if defined(__GLASGOW_HASKELL__) || defined(__MHS__) - 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 - --- | @since 0.5.9 -instance Read1 IntMap where - liftReadsPrec rp rl = readsData $ - readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList - where - rp' = liftReadsPrec rp rl - rl' = liftReadList rp rl - -{-------------------------------------------------------------------- - Helpers ---------------------------------------------------------------------} -{-------------------------------------------------------------------- - Link ---------------------------------------------------------------------} - --- | Link two @IntMap@s. The maps must not be empty. The @Prefix@es of the two --- maps must be different. @k1@ must share the prefix of @t1@. @p2@ must be the --- prefix of @t2@. -linkKey :: Key -> IntMap a -> Prefix -> IntMap a -> IntMap a -linkKey k1 t1 p2 t2 = link k1 t1 (unPrefix p2) t2 -{-# INLINE linkKey #-} - --- | Link two @IntMap@s. The maps must not be empty. The @Prefix@es of the two --- maps must be different. @k1@ must share the prefix of @t1@ and @k2@ must --- share the prefix of @t2@. -link :: Int -> IntMap a -> Int -> IntMap a -> IntMap a -link k1 t1 k2 t2 = linkWithMask (branchMask k1 k2) k1 t1 k2 t2 -{-# INLINE link #-} - --- `linkWithMask` is useful when the `branchMask` has already been computed -linkWithMask :: Int -> Key -> IntMap a -> Key -> IntMap a -> IntMap a -linkWithMask m k1 t1 k2 t2 - | i2w k1 < i2w k2 = Bin p t1 t2 - | otherwise = Bin p t2 t1 - where - p = Prefix (mask k1 m .|. m) -{-# INLINE linkWithMask #-} - -{-------------------------------------------------------------------- - @bin@ assures that we never have empty trees within a tree. ---------------------------------------------------------------------} - -bin :: Prefix -> IntMap a -> IntMap a -> IntMap a -bin _ l Nil = l -bin _ Nil r = r -bin p l r = Bin p l r -{-# INLINE bin #-} - --- binCheckLeft only checks that the left subtree is non-empty -binCheckLeft :: Prefix -> IntMap a -> IntMap a -> IntMap a -binCheckLeft _ Nil r = r -binCheckLeft p l r = Bin p l r -{-# INLINE binCheckLeft #-} - --- binCheckRight only checks that the right subtree is non-empty -binCheckRight :: Prefix -> IntMap a -> IntMap a -> IntMap a -binCheckRight _ l Nil = l -binCheckRight p l r = Bin p l r -{-# INLINE binCheckRight #-} - -{-------------------------------------------------------------------- - Utilities ---------------------------------------------------------------------} - --- | \(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 p l r - | signBranch p -> [r, l] - | otherwise -> [l, r] -{-# INLINE splitRoot #-} - - -{-------------------------------------------------------------------- - Debugging ---------------------------------------------------------------------} - --- | \(O(n \min(n,W))\). 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 \min(n,W))\). 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 l r -> - showsTree wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . showString (showBin p) . 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 l r -> - showsBars bars . showString (showBin p) . 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 -> 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 - _ : tl -> showString (concat (reverse tl)) . showString node - -node :: String -node = "+--" - -withBar, withEmpty :: [String] -> [String] -withBar bars = "| ":bars -withEmpty bars = " ":bars - -{-------------------------------------------------------------------- - Notes ---------------------------------------------------------------------} - --- Note [Okasaki-Gill] --- ~~~~~~~~~~~~~~~~~~~ --- --- The IntMap structure is based on the map described in the paper "Fast --- Mergeable Integer Maps" by Chris Okasaki and Andy Gill, with some --- differences. --- --- The paper spends most of its time describing a little-endian tree, where the --- branching is done first on low bits then high bits. It then briefly describes --- a big-endian tree. The implementation here is big-endian. --- --- The definition of Okasaki and Gill's map would be written in Haskell as --- --- data Dict a --- = Empty --- | Lf !Int a --- | Br !Int !Int !(Dict a) !(Dict a) --- --- Empty is the same as IntMap's Nil, and Lf is the same as Tip. --- --- In Br, the first Int is the shared prefix and the second is the mask bit by --- itself. For the big-endian map, the paper suggests that the prefix be the --- common prefix, followed by a 0-bit, followed by all 1-bits. This is so that --- the prefix value can be used as a point of split for binary search. --- --- IntMap's Bin corresponds to Br, but is different because it has only one --- Int (newtyped as Prefix). This describes both prefix and mask, so it is not --- necessary to store them separately. This value is, in fact, one plus the --- value suggested for the prefix in the paper. This representation is chosen --- because it saves one word per Bin without detriment to the efficiency of --- operations. --- --- The implementation of operations such as lookup, insert, union, follow --- the described implementations on Dict and split into the same cases. For --- instance, for insert, the three cases on a Br are whether the key belongs --- outside the map, or it belongs in the left child, or it belongs in the --- right child. We have the same three cases for a Bin. However, the bitwise --- operations we use to determine the case is naturally different due to the --- difference in representation. - --- Note [IntMap merge complexity] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- The merge algorithm (used for union, intersection, etc.) is adopted from --- Okasaki-Gill who give the complexity as O(m+n), where m and n are the sizes --- of the two input maps. This is correct, since we visit all constructors in --- both maps in the worst case, but we can try to find a tighter bound. --- --- Consider that m<=n, i.e. m is the size of the smaller map and n is the size --- of the larger. It does not matter which map is the first argument. --- --- Now we have O(n) as one upper bound for our complexity, since O(n) is the --- same as O(m+n) for m<=n. --- --- Next, consider the smaller map. For this map, we will visit some --- constructors, plus all the Bins of the larger map that lie in our way. --- For the former, the worst case is that we visit all constructors, which is --- O(m). --- For the latter, the worst case is that we encounter Bins at every point --- possible. This happens when for every key in the smaller map, the path to --- that key's Tip in the larger map has a full length of W, with a Bin at every --- bit position. To maximize the total number of Bins, the paths should be as --- disjoint as possible. But even if the paths are spread out, at least O(m) --- Bins are unavoidably shared, which extend up to a depth of lg(m) from the --- root. Beyond this, the paths may be disjoint. This gives us a total of --- O(m + m (W - lg m)) = O(m log (2^W / m)). --- The number of Bins we encounter is also bounded by the total number of Bins, --- which is n-1, but we already have O(n) as an upper bound. --- --- Combining our bounds, we have the final complexity as --- O(min(n, m log (2^W / m))). --- --- Note that --- * This is similar to the Map merge complexity, which is O(m log (n/m)). --- * When m is a small constant the term simplifies to O(min(n, W)), which is --- just the complexity we expect for single operations like insert and delete. +#ifdef __GLASGOW_HASKELL__ +{-# LANGUAGE DeriveLift #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE Trustworthy #-} +#endif + +{-# OPTIONS_HADDOCK not-home #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntMap.Internal +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +-- (c) wren romano 2016 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Portability : portable +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- The contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- +-- = Finite Int Maps (lazy interface internals) +-- +-- The @'IntMap' v@ type represents a finite map (sometimes called a dictionary) +-- from keys of type @Int@ to values of type @v@. +-- +-- +-- == Implementation +-- +-- The implementation is based on /big-endian patricia trees/. This data +-- structure performs especially well on binary operations like 'union' +-- and 'intersection'. Additionally, 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, +-- . +-- +-- @since 0.5.9 +----------------------------------------------------------------------------- + +-- [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.Internal ( + -- * Map type + IntMap(..) -- instance Eq,Show + , Key + + -- * Operators + , (!), (!?), (\\) + + -- * Query + , null + , size + , compareSize + , member + , notMember + , lookup + , findWithDefault + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , disjoint + + -- * Construction + , empty + , singleton + , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA + + -- ** Insertion + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + + -- ** Delete\/Update + , delete + , pop + , adjust + , adjustWithKey + , update + , updateWithKey + , upsert + , updateLookupWithKey + , alter + , alterF + + -- * Combine + + -- ** Union + , union + , unionWith + , unionWithKey + , unions + , unionsWith + + -- ** Difference + , difference + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey + + -- ** Symmetric difference + , symmetricDifference + + -- ** Compose + , compose + + -- ** General combining function + , SimpleWhenMissing + , SimpleWhenMatched + , runWhenMatched + , runWhenMissing + , merge + -- *** @WhenMatched@ tactics + , dropMatched + , zipWithMaybeMatched + , zipWithMatched + -- *** @WhenMissing@ tactics + , mapMaybeMissing + , dropMissing + , preserveMissing + , mapMissing + , filterMissing + + -- ** Applicative general combining function + , WhenMissing (..) + , WhenMatched (..) + , mergeA + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , zipWithMaybeAMatched + , zipWithAMatched + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for + -- 'mergeA' as well. Furthermore, the following + -- are available. + , traverseMaybeMissing + , traverseMissing + , filterAMissing + , whenMissing + + -- ** Deprecated general combining function + , mergeWithKey + , mergeWithKey' + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey + , traverseMaybeWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldl' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + + -- ** Lists + , toList + , fromList + , fromListWith + , fromListWithKey + , fromListUpsert + + -- ** Ordered lists + , toAscList + , toDescList + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromAscListUpsert + , fromDistinctAscList + , fromDescList + , fromDescListUpsert + + -- * Filter + , filter + , filterKeys + , filterWithKey + , restrictKeys + , withoutKeys + , partition + , partitionWithKey + + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Min\/Max + , lookupMin + , lookupMax + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , updateMinWithKey + , updateMaxWithKey + , minView + , maxView + , minViewWithKey + , maxViewWithKey + + -- * Debugging + , showTree + , showTreeWith + + -- * Utility + , link + , linkKey + , bin + , binCheckL + , binCheckR + , MonoState(..) + , Stack(..) + , ascLinkTop + , ascLinkAll + , descInsert + , descLinkTop + , descLinkAll + , IntMapBuilder(..) + , BStack(..) + , emptyB + , insertB + , finishB + , moveToB + , MoveResult(..) + , treeFromIntSetTip + + -- * Used by "IntMap.Merge.Lazy" and "IntMap.Merge.Strict" + , mapWhenMissing + , mapWhenMatched + , lmapWhenMissing + , contramapFirstWhenMatched + , contramapSecondWhenMatched + , mapGentlyWhenMissing + , mapGentlyWhenMatched + ) where + +import Data.Functor.Identity (Identity (..)) +import Data.Semigroup (Semigroup(stimes)) +#if !(MIN_VERSION_base(4,11,0)) +import Data.Semigroup (Semigroup((<>))) +#endif +import Data.Semigroup (stimesIdempotentMonoid) +import Data.Functor.Classes + +import Control.DeepSeq (NFData(rnf),NFData1(liftRnf)) +import Data.Bits +import qualified Data.Foldable as Foldable +import Data.Maybe (fromMaybe) +import Utils.Containers.Internal.Prelude hiding + (lookup, map, filter, foldr, foldl, foldl', foldMap, null) +import Prelude () + +import Data.IntSet.Internal (IntSet) +import qualified Data.IntSet.Internal as IntSet +import Data.IntSet.Internal.IntTreeCommons + ( Key + , Prefix(..) + , nomatch + , left + , signBranch + , mask + , branchMask + , branchPrefix + , TreeTreeBranch(..) + , treeTreeBranch + , i2w + , Order(..) + ) +import Utils.Containers.Internal.BitUtil (shiftLL, shiftRL, iShiftRL, wordSize) +import Utils.Containers.Internal.Strict + (StrictPair(..), StrictTriple(..), toPair) + +#ifdef __GLASGOW_HASKELL__ +import Data.Coerce +import Data.Data (Data(..), Constr, mkConstr, constrIndex, + DataType, mkDataType, gcast1) +import qualified Data.Data as Data +import GHC.Exts (build) +import qualified GHC.Exts as GHCExts +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else +import Language.Haskell.TH.Syntax (Lift) +-- See Note [ Template Haskell Dependencies ] +import Language.Haskell.TH () +# endif +#endif +#if defined(__GLASGOW_HASKELL__) || defined(__MHS__) +import Text.Read +#endif +import qualified Control.Category as Category + + +{-------------------------------------------------------------------- + Types +--------------------------------------------------------------------} + + +-- | A map of integers to values @a@. + +-- See Note: Order of constructors +data IntMap a = Bin {-# UNPACK #-} !Prefix + !(IntMap a) + !(IntMap a) + | Tip {-# UNPACK #-} !Key a + | Nil + +-- +-- Note [IntMap structure and invariants] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- +-- * Nil is never found as a child of Bin. +-- +-- * The Prefix of a Bin indicates the common high-order bits that all keys in +-- the Bin share. +-- +-- * The least significant set bit of the Int value of a Prefix is called the +-- mask bit. +-- +-- * All the bits to the left of the mask bit are called the shared prefix. All +-- keys stored in the Bin begin with the shared prefix. +-- +-- * All keys in the left child of the Bin have the mask bit unset, and all keys +-- in the right child have the mask bit set. It follows that +-- +-- 1. The Int value of the Prefix of a Bin is the smallest key that can be +-- present in the right child of the Bin. +-- +-- 2. All keys in the right child of a Bin are greater than keys in the +-- left child, with one exceptional situation. If the Bin separates +-- negative and non-negative keys, the mask bit is the sign bit and the +-- left child stores the non-negative keys while the right child stores the +-- negative keys. +-- +-- * All bits to the right of the mask bit are set to 0 in a Prefix. +-- +-- +-- As an example, consider that on a 32-bit system we have a Bin with Prefix +-- +-- 0b00000000000100100000011000000000 +-- ^ mask bit +-- ^^^^^^^^^^^^^^^^^^^^^^ shared prefix +-- +-- The key +-- 0b00000000000100100000010000010100 belongs under this Bin, since it matches +-- the shared prefix. The mask bit is 0, so it belongs in the left child and not +-- the right. +-- +-- The key +-- 0b00000000000100000000010000010100 does not belong under this Bin, since it +-- does not match the shared prefix. + +-- See Note [Okasaki-Gill] for how the implementation here relates to the one in +-- Okasaki and Gill's paper. + +#ifdef __GLASGOW_HASKELL__ +-- | @since 0.6.6 +deriving instance Lift a => Lift (IntMap a) +#endif + +{-------------------------------------------------------------------- + Operators +--------------------------------------------------------------------} + +-- | \(O(\min(n,W))\). Find the value at a key. +-- Calls 'error' when the element can not be found. +-- +-- __Note__: This function is partial. Prefer '!?'. +-- +-- > 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 + +-- | \(O(\min(n,W))\). Find the value at a key. +-- Returns 'Nothing' when the element can not be found. +-- +-- > fromList [(5,'a'), (3,'b')] !? 1 == Nothing +-- > fromList [(5,'a'), (3,'b')] !? 5 == Just 'a' +-- +-- @since 0.5.11 + +(!?) :: IntMap a -> Key -> Maybe a +(!?) m k = lookup 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 +--------------------------------------------------------------------} + +-- | @mempty@ = 'empty' +instance Monoid (IntMap a) where + mempty = empty + mconcat = unions +#if !MIN_VERSION_base(4,11,0) + mappend = (<>) +#endif + +-- | @(<>)@ = 'union' +-- +-- @since 0.5.7 +instance Semigroup (IntMap a) where + (<>) = union + stimes = stimesIdempotentMonoid + +-- | Folds in order of increasing key. +instance Foldable.Foldable IntMap where + fold = foldMap id + {-# INLINABLE fold #-} + foldr = foldr + {-# INLINE foldr #-} + foldl = foldl + {-# INLINE foldl #-} + foldMap = foldMap + {-# INLINE foldMap #-} + foldl' = foldl' + {-# INLINE foldl' #-} + foldr' = foldr' + {-# INLINE foldr' #-} + length = size + {-# INLINE length #-} + null = null + {-# INLINE null #-} + toList = elems -- NB: Foldable.toList /= IntMap.toList + {-# INLINE toList #-} + elem = go + where go !_ Nil = False + go x (Tip _ y) = x == y + go x (Bin _ l r) = go x l || go x r + {-# INLINABLE elem #-} + maximum = start + where start Nil = error "Data.Foldable.maximum (for Data.IntMap): empty map" + start (Tip _ y) = y + start (Bin p l r) + | signBranch p = go (start r) l + | otherwise = go (start l) r + + go !m Nil = m + go m (Tip _ y) = max m y + go m (Bin _ l r) = go (go m l) r + {-# INLINABLE maximum #-} + minimum = start + where start Nil = error "Data.Foldable.minimum (for Data.IntMap): empty map" + start (Tip _ y) = y + start (Bin p l r) + | signBranch p = go (start r) l + | otherwise = go (start l) r + + go !m Nil = m + go m (Tip _ y) = min m y + go m (Bin _ l r) = go (go m l) r + {-# INLINABLE minimum #-} + sum = foldl' (+) 0 + {-# INLINABLE sum #-} + product = foldl' (*) 1 + {-# INLINABLE product #-} + +-- | Traverses in order of increasing key. +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 + +-- | @since 0.8 +instance NFData1 IntMap where + liftRnf rnfx = go + where + go Nil = () + go (Tip _ v) = rnfx v + go (Bin _ l r) = go l `seq` go 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" [] Data.Prefix + +intMapDataType :: DataType +intMapDataType = mkDataType "Data.IntMap.Internal.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 entries in the map. +-- +-- __Note__: Unlike @Data.Map.'Data.Map.Lazy.size'@, this is /not/ \(O(1)\). +-- +-- > size empty == 0 +-- > size (singleton 1 'a') == 1 +-- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 +-- +-- See also: 'compareSize' +size :: IntMap a -> Int +size = go 0 + where + go !acc (Bin _ l r) = go (go acc l) r + go acc (Tip _ _) = 1 + acc + go acc Nil = acc + +-- | \(O(\min(n,c))\). Compare the number of entries in the map to an @Int@. +-- +-- @compareSize m c@ returns the same result as @compare ('size' m) c@ but is +-- more efficient when @c@ is smaller than the size of the map. +-- +-- @since 0.8.1 +compareSize :: IntMap a -> Int -> Ordering +compareSize Nil c0 = compare 0 c0 +compareSize _ c0 | c0 <= 0 = GT +compareSize t c0 = compare 0 (go t (c0 - 1)) + where + go (Bin _ _ _) 0 = -1 + go (Bin _ l r) c + | c' < 0 = c' + | otherwise = go r c' + where + c' = go l (c - 1) + go _ c = c + +-- | \(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 = go + where + go (Bin p l r) + | nomatch k p = False + | left k p = go l + | otherwise = go r + go (Tip kx _) = k == kx + go Nil = False + +-- | \(O(\min(n,W))\). Is the key not a member of the map? +-- +-- > notMember 5 (fromList [(5,'a'), (3,'b')]) == False +-- > notMember 1 (fromList [(5,'a'), (3,'b')]) == True + +notMember :: Key -> IntMap a -> Bool +notMember k m = not $ member k m + +-- | \(O(\min(n,W))\). Look up 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 = go + where + go (Bin p l r) | left k p = 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 = go + where + go (Bin p l r) | left k p = 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 = go + where + go (Bin p l r) | nomatch k p = def + | left k p = go l + | otherwise = go r + go (Tip kx x) | k == kx = x + | otherwise = def + go Nil = def + +-- | \(O(\min(n,W))\). 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 = case t of + Bin p l r | signBranch p -> if k >= 0 then go r l else go Nil r + _ -> go Nil t + where + go def (Bin p l r) + | nomatch k p = if k < unPrefix p then unsafeFindMax def else unsafeFindMax r + | left k p = 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(\min(n,W))\). 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 = case t of + Bin p l r | signBranch p -> if k >= 0 then go Nil l else go l r + _ -> go Nil t + where + go def (Bin p l r) + | nomatch k p = if k < unPrefix p then unsafeFindMin l else unsafeFindMin def + | left k p = 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(\min(n,W))\). 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 = case t of + Bin p l r | signBranch p -> if k >= 0 then go r l else go Nil r + _ -> go Nil t + where + go def (Bin p l r) + | nomatch k p = if k < unPrefix p then unsafeFindMax def else unsafeFindMax r + | left k p = 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(\min(n,W))\). 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 = case t of + Bin p l r | signBranch p -> if k >= 0 then go Nil l else go l r + _ -> go Nil t + where + go def (Bin p l r) + | nomatch k p = if k < unPrefix p then unsafeFindMin l else unsafeFindMin def + | left k p = 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 + +{-------------------------------------------------------------------- + Disjoint +--------------------------------------------------------------------} +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- Check whether the key sets of two maps are disjoint +-- (i.e. their 'intersection' is empty). +-- +-- > disjoint (fromList [(2,'a')]) (fromList [(1,()), (3,())]) == True +-- > disjoint (fromList [(2,'a')]) (fromList [(1,'a'), (2,'b')]) == False +-- > disjoint (fromList []) (fromList []) == True +-- +-- > disjoint a b == null (intersection a b) +-- +-- @since 0.6.2.1 +disjoint :: IntMap a -> IntMap b -> Bool +disjoint Nil _ = True +disjoint _ Nil = True +disjoint (Tip kx _) ys = notMember kx ys +disjoint xs (Tip ky _) = notMember ky xs +disjoint t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> disjoint l1 t2 + ABR -> disjoint r1 t2 + BAL -> disjoint t1 l2 + BAR -> disjoint t1 r2 + EQL -> disjoint l1 l2 && disjoint r1 r2 + NOM -> True + +{-------------------------------------------------------------------- + Compose +--------------------------------------------------------------------} +-- | Given maps @bc@ and @ab@, relate the keys of @ab@ to the values of @bc@, +-- by using the values of @ab@ as keys for lookups in @bc@. +-- +-- Complexity: \( O(n * \min(m,W)) \), where \(m\) is the size of the first argument +-- +-- > compose (fromList [('a', "A"), ('b', "B")]) (fromList [(1,'a'),(2,'b'),(3,'z')]) = fromList [(1,"A"),(2,"B")] +-- +-- @ +-- ('compose' bc ab '!?') = (bc '!?') <=< (ab '!?') +-- @ +-- +-- __Note:__ Prior to v0.6.4, "Data.IntMap.Strict" exposed a version of +-- 'compose' that forced the values of the output 'IntMap'. This version does +-- not force these values. +-- +-- @since 0.6.3.1 +compose :: IntMap c -> IntMap Int -> IntMap c +compose bc !ab + | null bc = empty + | otherwise = mapMaybe (bc !?) ab + +{-------------------------------------------------------------------- + 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@(Bin p l r) + | nomatch k p = linkKey k (Tip k x) p t + | left k p = Bin p (insert k x l) r + | otherwise = Bin p l (insert k x r) +insert k x t@(Tip ky _) + | k==ky = Tip k x + | otherwise = link k (Tip k x) ky t +insert k x Nil = Tip k x + +-- right-biased insertion, used by 'union' +-- | \(O(\min(n,W))\). Insert with a combining function. +-- @'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" +-- +-- Also see the performance note on 'fromListWith'. + +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" +-- +-- Also see the performance note on 'fromListWith'. + +insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithKey f !k x t@(Bin p l r) + | nomatch k p = linkKey k (Tip k x) p t + | left k p = Bin p (insertWithKey f k x l) r + | otherwise = Bin p l (insertWithKey f k x r) +insertWithKey f k x t@(Tip ky y) + | k == ky = Tip k (f k x y) + | otherwise = link k (Tip k x) ky t +insertWithKey _ k x Nil = Tip k x + +-- | \(O(\min(n,W))\). The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- 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")]) +-- +-- Also see the performance note on 'fromListWith'. + +insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) +insertLookupWithKey f !k x t@(Bin p l r) + | nomatch k p = (Nothing,linkKey k (Tip k x) p t) + | left k p = let (found,l') = insertLookupWithKey f k x l + in (found,Bin p l' r) + | otherwise = let (found,r') = insertLookupWithKey f k x r + in (found,Bin p l r') +insertLookupWithKey f k x t@(Tip ky y) + | k == ky = (Just y,Tip k (f k x y)) + | otherwise = (Nothing,link k (Tip k x) ky t) +insertLookupWithKey _ k x Nil = (Nothing,Tip k x) + + +{-------------------------------------------------------------------- + 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@(Bin p l r) + | nomatch k p = t + | left k p = binCheckL p (delete k l) r + | otherwise = binCheckR p l (delete k r) +delete k t@(Tip ky _) + | k == ky = Nil + | otherwise = t +delete _k Nil = Nil + +-- | \(O(\min(n,W))\). Pop an entry from the map. +-- +-- Returns @Nothing@ if the key is not in the map. Otherwise returns @Just@ the +-- value at the key and a map with the entry removed. +-- +-- @ +-- pop 1 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Nothing +-- pop 2 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Just ("b",fromList [(0,"a"),(4,"c")]) +-- @ +-- +-- @since 0.8.1 +pop :: Key -> IntMap a -> Maybe (a, IntMap a) +pop k0 t0 = case go k0 t0 of + Popped (Just y) t -> Just (y, t) + _ -> Nothing + where + go !k (Bin p l r) + | nomatch k p = Popped Nothing Nil + | left k p = case go k l of + Popped y@(Just _) l' -> Popped y (binCheckL p l' r) + q -> q + | otherwise = case go k r of + Popped y@(Just _) r' -> Popped y (binCheckR p l r') + q -> q + go !k (Tip ky y) + | k == ky = Popped (Just y) Nil + | otherwise = Popped Nothing Nil + go !_ Nil = Popped Nothing Nil + +-- See Note [Popped impl] in Data.Map.Internal +data Popped k a = Popped +#if __GLASGOW_HASKELL__ >= 906 + {-# UNPACK #-} +#endif + !(Maybe a) + !(IntMap a) + +-- | \(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 !k (Bin p l r) + | left k p = Bin p (adjustWithKey f k l) r + | otherwise = Bin p l (adjustWithKey f k r) +adjustWithKey f k t@(Tip ky y) + | k == ky = Tip ky (f k y) + | otherwise = t +adjustWithKey _ _ Nil = Nil + + +-- | \(O(\min(n,W))\). The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > 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 (Bin p l r) + | left k p = binCheckL p (updateWithKey f k l) r + | otherwise = binCheckR p l (updateWithKey f k r) +updateWithKey f k t@(Tip ky y) + | k == ky = case (f k y) of + Just y' -> Tip ky y' + Nothing -> Nil + | otherwise = t +updateWithKey _ _ Nil = Nil + +-- | \(O(\min(n,W))\). Update the value at a key or insert a value if the key is +-- not in the map. +-- +-- @ +-- let inc = maybe 1 (+1) +-- upsert inc 100 (fromList [(100,1),(300,2)]) == fromList [(100,2),(300,2)] +-- upsert inc 200 (fromList [(100,1),(300,2)]) == fromList [(100,1),(200,1),(300,2)] +-- @ +-- +-- @since 0.8.1 +upsert :: (Maybe a -> a) -> Key -> IntMap a -> IntMap a +upsert f !k t@(Bin p l r) + | nomatch k p = linkKey k (Tip k (f Nothing)) p t + | left k p = Bin p (upsert f k l) r + | otherwise = Bin p l (upsert f k r) +upsert f !k t@(Tip ky y) + | k == ky = Tip ky (f (Just y)) + | otherwise = link k (Tip k (f Nothing)) ky t +upsert f !k Nil = Tip k (f Nothing) + +-- | \(O(\min(n,W))\). Look up and update. +-- This function returns the 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") +-- +-- See also: 'pop' + +updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a) +updateLookupWithKey f !k (Bin p l r) + | left k p = let !(found,l') = updateLookupWithKey f k l + in (found,binCheckL p l' r) + | otherwise = let !(found,r') = updateLookupWithKey f k r + in (found,binCheckR p l r') +updateLookupWithKey f k t@(Tip ky y) + | k==ky = case (f k y) of + Just y' -> (Just y,Tip ky y') + Nothing -> (Just y,Nil) + | otherwise = (Nothing,t) +updateLookupWithKey _ _ Nil = (Nothing,Nil) + + + +-- | \(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@(Bin p l r) + | nomatch k p = case f Nothing of + Nothing -> t + Just x -> linkKey k (Tip k x) p t + | left k p = binCheckL p (alter f k l) r + | otherwise = binCheckR p l (alter f k r) +alter f k t@(Tip ky y) + | k==ky = case f (Just y) of + Just x -> Tip ky x + Nothing -> Nil + | otherwise = case f Nothing of + Just x -> link k (Tip k x) ky t + Nothing -> Tip ky y +alter f k Nil = case f Nothing of + Just x -> Tip k x + Nothing -> Nil + +-- | \(O(\min(n,W))\). The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, +-- or update a value in an 'IntMap'. In short : @'lookup' k \<$\> 'alterF' f k m = f +-- ('lookup' k m)@. +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. +-- +-- Note: 'alterF' is a flipped version of the @at@ combinator from +-- @Control.Lens.At@. +-- +-- === Examples +-- +-- @ +-- -- Lookup the value at the key, and also remove the existing value or set a new value. +-- lookupAndSet :: Key -> Maybe a -> IntMap a -> (Maybe a, IntMap a) +-- lookupAndSet k new = alterF (\\old -> (old, new)) k +-- @ +-- +-- @ +-- -- Delete the value at the key. If it is absent the result is Nothing. +-- mustDelete :: Key -> IntMap a -> Maybe (IntMap a) +-- mustDelete = alterF (Nothing <$) +-- @ +-- +-- @since 0.5.8 + +alterF :: Functor f + => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) +-- This implementation was stolen from 'Control.Lens.At'. +alterF f k m = (<$> f mv) $ \fres -> + case fres of + Nothing -> maybe m (const (delete k m)) mv + Just v' -> insert k v' m + where mv = lookup k m +{-# INLINE alterF #-} + +{-------------------------------------------------------------------- + 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 :: Foldable f => f (IntMap a) -> IntMap a +unions xs + = Foldable.foldl' union empty xs +{-# INLINE unions #-} -- Inline for list fusion + +-- | 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 :: Foldable f => (a->a->a) -> f (IntMap a) -> IntMap a +unionsWith f ts + = Foldable.foldl' (unionWith f) empty ts +{-# INLINE unionsWith #-} -- Inline for list fusion + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- The (left-biased) union of two maps. +-- It prefers the first map when duplicate keys are encountered, +-- i.e. (@'union' == 'unionWith' 'const'@). +-- +-- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] + +union :: IntMap a -> IntMap a -> IntMap a +union m1 m2 + = mergeWithKey' Bin const id id m1 m2 + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- The union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] +-- +-- Also see the performance note on 'fromListWith'. + +unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWith f = unionWithKey (\_ x y -> f x y) +{-# INLINE unionWith #-} + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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")] +-- +-- Also see the performance note on 'fromListWith'. + +unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithKey f m1 m2 + = mergeWithKey' Bin f' id id m1 m2 + where + f' (Tip k1 x1) (Tip _k2 x2) = Tip k1 (f k1 x1 x2) + f' _ _ = error "not Tip" +{-# INLINABLE unionWithKey #-} -- See Note [INLINABLE to expose unfoldings] + +{-------------------------------------------------------------------- + Difference +--------------------------------------------------------------------} +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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 = differenceWithKey (\_ x y -> f x y) +{-# INLINE differenceWith #-} + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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 +{-# INLINABLE differenceWithKey #-} -- See Note [INLINABLE to expose unfoldings] + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- Remove all the keys in a given set from a map. +-- +-- @ +-- m \`withoutKeys\` s = 'filterWithKey' (\\k _ -> k ``IntSet.notMember`` s) m +-- @ +-- +-- @since 0.5.8 +withoutKeys :: IntMap a -> IntSet -> IntMap a +withoutKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> binCheckL p1 (withoutKeys l1 t2) r1 + ABR -> binCheckR p1 l1 (withoutKeys r1 t2) + BAL -> withoutKeys t1 l2 + BAR -> withoutKeys t1 r2 + EQL -> bin p1 (withoutKeys l1 l2) (withoutKeys r1 r2) + NOM -> t1 + where +withoutKeys t1@(Bin _ _ _) (IntSet.Tip p2 bm2) = withoutKeysTip t1 p2 bm2 +withoutKeys t1@(Bin _ _ _) IntSet.Nil = t1 +withoutKeys t1@(Tip k1 _) t2 + | k1 `IntSet.member` t2 = Nil + | otherwise = t1 +withoutKeys Nil _ = Nil + +withoutKeysTip :: IntMap a -> Int -> IntSet.BitMap -> IntMap a +withoutKeysTip t@(Bin p l r) !p2 !bm2 + | IntSet.suffixOf (unPrefix p) /= 0 = + if IntSet.prefixOf (unPrefix p) == p2 + then restrictBM t (complement bm2) + else t + | nomatch p2 p = t + | left p2 p = binCheckL p (withoutKeysTip l p2 bm2) r + | otherwise = binCheckR p l (withoutKeysTip r p2 bm2) +withoutKeysTip t@(Tip kx _) !p2 !bm2 + | IntSet.prefixOf kx == p2 && IntSet.bitmapOf kx .&. bm2 /= 0 = Nil + | otherwise = t +withoutKeysTip Nil !_ !_ = Nil + +{-------------------------------------------------------------------- + Intersection +--------------------------------------------------------------------} +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- The (left-biased) intersection of two maps (based on keys). +-- +-- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" + +intersection :: IntMap a -> IntMap b -> IntMap a +intersection m1 m2 + = mergeWithKey' bin const (const Nil) (const Nil) m1 m2 + + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- The restriction of a map to the keys in a set. +-- +-- @ +-- m \`restrictKeys\` s = 'filterWithKey' (\\k _ -> k ``IntSet.member`` s) m +-- @ +-- +-- @since 0.5.8 +restrictKeys :: IntMap a -> IntSet -> IntMap a +restrictKeys t1@(Bin p1 l1 r1) t2@(IntSet.Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> restrictKeys l1 t2 + ABR -> restrictKeys r1 t2 + BAL -> restrictKeys t1 l2 + BAR -> restrictKeys t1 r2 + EQL -> bin p1 (restrictKeys l1 l2) (restrictKeys r1 r2) + NOM -> Nil +restrictKeys t1@(Bin _ _ _) (IntSet.Tip p2 bm2) = restrictKeysTip t1 p2 bm2 +restrictKeys (Bin _ _ _) IntSet.Nil = Nil +restrictKeys t1@(Tip k1 _) t2 + | k1 `IntSet.member` t2 = t1 + | otherwise = Nil +restrictKeys Nil _ = Nil + +restrictKeysTip :: IntMap a -> Int -> IntSet.BitMap -> IntMap a +restrictKeysTip t@(Bin p l r) !p2 !bm2 + | IntSet.suffixOf (unPrefix p) /= 0 = + if IntSet.prefixOf (unPrefix p) == p2 + then restrictBM t bm2 + else Nil + | nomatch p2 p = Nil + | left p2 p = restrictKeysTip l p2 bm2 + | otherwise = restrictKeysTip r p2 bm2 +restrictKeysTip t@(Tip kx _) !p2 !bm2 + | IntSet.prefixOf kx == p2 && IntSet.bitmapOf kx .&. bm2 /= 0 = t + | otherwise = Nil +restrictKeysTip Nil !_ !_ = Nil + +-- Must be called on an IntMap whose keys fit in the given IntSet BitMap's Tip. +-- Keeps keys that match the BitMap. +-- Returns early as an optimization, i.e. if the tree can be entirely kept or +-- discarded there is no need to recursively visit the children. +restrictBM :: IntMap a -> IntSet.BitMap -> IntMap a +restrictBM t@(Bin p l r) !bm + | bm' == 0 = Nil + | bm' == -1 = t + | otherwise = bin p (restrictBM l bm) (restrictBM r bm) + where + -- Here we care about the "submask" of bm corresponding the current Bin's + -- range. So we create bm', where this submask is at the lowest position and + -- and all other bits are set to the highest bit of the submask (using an + -- arithmetic shiftR). Now bm' is 0 when the submask is empty and -1 when + -- the submask is full. + px = IntSet.suffixOf (unPrefix p) + px1 = px - 1 + min_ = px .&. px1 + max_ = px .|. px1 + sh = (wordSize - 1) - max_ + bm' = (w2i bm `unsafeShiftL` sh) `unsafeShiftR` (sh + min_) +restrictBM t@(Tip k _) !bm + | IntSet.bitmapOf k .&. bm /= 0 = t + | otherwise = Nil +restrictBM Nil !_ = Nil + +w2i :: Word -> Int +w2i = fromIntegral + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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 = intersectionWithKey (\_ x y -> f x y) +{-# INLINE intersectionWith #-} + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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 f' (const Nil) (const Nil) m1 m2 + where + f' (Tip k1 x1) (Tip _k2 x2) = Tip k1 (f k1 x1 x2) + f' _ _ = error "not Tip" +-- See Note [INLINABLE to expose unfoldings] +{-# INLINABLE intersectionWithKey #-} + +{-------------------------------------------------------------------- + Symmetric difference +--------------------------------------------------------------------} + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- The symmetric difference of two maps. +-- +-- The result contains entries whose keys appear in exactly one of the two maps. +-- +-- @ +-- symmetricDifference +-- (fromList [(0,\'q\'),(2,\'b\'),(4,\'w\'),(6,\'o\')]) +-- (fromList [(0,\'e\'),(3,\'r\'),(6,\'t\'),(9,\'s\')]) +-- == +-- fromList [(2,\'b\'),(3,\'r\'),(4,\'w\'),(9,\'s\')] +-- @ +-- +-- @since 0.8 +symmetricDifference :: IntMap a -> IntMap a -> IntMap a +symmetricDifference t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = + case treeTreeBranch p1 p2 of + ABL -> binCheckL p1 (symmetricDifference l1 t2) r1 + ABR -> binCheckR p1 l1 (symmetricDifference r1 t2) + BAL -> binCheckL p2 (symmetricDifference t1 l2) r2 + BAR -> binCheckR p2 l2 (symmetricDifference t1 r2) + EQL -> bin p1 (symmetricDifference l1 l2) (symmetricDifference r1 r2) + NOM -> link (unPrefix p1) t1 (unPrefix p2) t2 +symmetricDifference t1@(Bin _ _ _) t2@(Tip k2 _) = symDiffTip t2 k2 t1 +symmetricDifference t1@(Bin _ _ _) Nil = t1 +symmetricDifference t1@(Tip k1 _) t2 = symDiffTip t1 k1 t2 +symmetricDifference Nil t2 = t2 + +symDiffTip :: IntMap a -> Int -> IntMap a -> IntMap a +symDiffTip !t1 !k1 = go + where + go t2@(Bin p2 l2 r2) + | nomatch k1 p2 = linkKey k1 t1 p2 t2 + | left k1 p2 = binCheckL p2 (go l2) r2 + | otherwise = binCheckR p2 l2 (go r2) + go t2@(Tip k2 _) + | k1 == k2 = Nil + | otherwise = link k1 t1 k2 t2 + go Nil = t1 + +{-------------------------------------------------------------------- + MergeWithKey +--------------------------------------------------------------------} + +-- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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'). +-- +-- __Warning__: 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@. + +-- See Note [IntMap merge complexity] +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 + combine (Tip k1 x1) (Tip _k2 x2) = + case f k1 x1 x2 of + Nothing -> Nil + Just x -> Tip k1 x + combine _ _ = error "not Tip" + {-# 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 -> 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 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> bin' p1 (go l1 t2) (g1 r1) + ABR -> bin' p1 (g1 l1) (go r1 t2) + BAL -> bin' p2 (go t1 l2) (g2 r2) + BAR -> bin' p2 (g2 l2) (go t1 r2) + EQL -> bin' p1 (go l1 l2) (go r1 r2) + NOM -> maybe_link (unPrefix p1) (g1 t1) (unPrefix p2) (g2 t2) + + go t1'@(Bin _ _ _) t2'@(Tip k2' _) = merge0 t2' k2' t1' + where + merge0 t2 k2 t1@(Bin p1 l1 r1) + | nomatch k2 p1 = maybe_link (unPrefix p1) (g1 t1) k2 (g2 t2) + | left k2 p1 = bin' p1 (merge0 t2 k2 l1) (g1 r1) + | otherwise = bin' p1 (g1 l1) (merge0 t2 k2 r1) + merge0 t2 k2 t1@(Tip k1 _) + | k1 == k2 = f t1 t2 + | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2) + merge0 t2 _ Nil = g2 t2 + + go t1@(Bin _ _ _) Nil = g1 t1 + + go t1'@(Tip k1' _) t2' = merge0 t1' k1' t2' + where + merge0 t1 k1 t2@(Bin p2 l2 r2) + | nomatch k1 p2 = maybe_link k1 (g1 t1) (unPrefix p2) (g2 t2) + | left k1 p2 = bin' p2 (merge0 t1 k1 l2) (g2 r2) + | otherwise = bin' p2 (g2 l2) (merge0 t1 k1 r2) + merge0 t1 k1 t2@(Tip k2 _) + | k1 == k2 = f t1 t2 + | otherwise = maybe_link k1 (g1 t1) k2 (g2 t2) + merge0 t1 _ Nil = g1 t1 + + go Nil Nil = Nil + + go Nil t2 = g2 t2 + + maybe_link _ Nil _ t2 = t2 + maybe_link _ t1 _ Nil = t1 + maybe_link k1 t1 k2 t2 = link k1 t1 k2 t2 + {-# INLINE maybe_link #-} +{-# INLINE mergeWithKey' #-} + + +{-------------------------------------------------------------------- + mergeA +--------------------------------------------------------------------} + +-- | A tactic for dealing with keys present in one map but not the +-- other in 'merge' or 'mergeA'. +-- +-- A tactic of type @WhenMissing f k x z@ is an abstract representation +-- of a function of type @Key -> x -> f (Maybe z)@. +-- +-- @since 0.5.9 + +data WhenMissing f x y = WhenMissing + { missingSubtree :: IntMap x -> f (IntMap y) + , missingKey :: Key -> x -> f (Maybe y)} + +-- | @since 0.5.9 +instance Monad f => Functor (WhenMissing f x) where + fmap = mapWhenMissing + {-# INLINE fmap #-} + + +-- | @since 0.5.9 +instance Monad f => Category.Category (WhenMissing f) + where + id = preserveMissing + f . g = + traverseMaybeMissing $ \ k x -> do + y <- missingKey g k x + case y of + Nothing -> pure Nothing + Just q -> missingKey f k q + {-# INLINE id #-} + {-# INLINE (.) #-} + + +-- | Equivalent to @ReaderT k (ReaderT x (MaybeT f))@. +-- +-- @since 0.5.9 +instance Monad f => Applicative (WhenMissing f x) where + pure x = mapMissing (\ _ _ -> x) + f <*> g = + traverseMaybeMissing $ \k x -> do + res1 <- missingKey f k x + case res1 of + Nothing -> pure Nothing + Just r -> (pure $!) . fmap r =<< missingKey g k x + {-# INLINE pure #-} + {-# INLINE (<*>) #-} + + +-- | Equivalent to @ReaderT k (ReaderT x (MaybeT f))@. +-- +-- @since 0.5.9 +instance Monad f => Monad (WhenMissing f x) where + m >>= f = + traverseMaybeMissing $ \k x -> do + res1 <- missingKey m k x + case res1 of + Nothing -> pure Nothing + Just r -> missingKey (f r) k x + {-# INLINE (>>=) #-} + +-- | Create a @WhenMissing@ from two functions. +-- +-- @whenMissing@ must be called with two functions @f@ and @g@ such that +-- @g = 'traverseMaybeWithKey' f@. @g@ may be a more efficient way of applying +-- @f@ to all key-value pairs in an @IntMap@. +-- +-- __Warning__: It is the caller's responsibility to ensure the above property. +-- +-- === __Examples__ +-- +-- @ +-- preserveMissing :: Applicative f => WhenMissing f x x +-- preserveMissing = whenMissing f g +-- where +-- f _k x = pure (Just x) +-- g m = pure m +-- -- Note that this satisfies g = traverseMaybeWithKey f +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- For a usage of this, see examples on mergeA +-- isEmpty :: WhenMissing (Const All) x y +-- isEmpty = whenMissing f g +-- where +-- f _k _x = Const (All False) +-- g m = Const (All (null m)) +-- -- Note that this satisfies g = traverseMaybeWithKey f +-- @ +-- +-- @since 0.8.1 +whenMissing + :: (Key -> x -> f (Maybe y)) + -> (IntMap x -> f (IntMap y)) + -> WhenMissing f x y +whenMissing = flip WhenMissing + +-- | Map covariantly over a @'WhenMissing' f x@. +-- +-- @since 0.5.9 +mapWhenMissing + :: Monad f + => (a -> b) + -> WhenMissing f x a + -> WhenMissing f x b +mapWhenMissing f t = WhenMissing + { missingSubtree = \m -> missingSubtree t m >>= \m' -> pure $! fmap f m' + , missingKey = \k x -> missingKey t k x >>= \q -> (pure $! fmap f q) } +{-# INLINE mapWhenMissing #-} + + +-- | Map covariantly over a @'WhenMissing' f x@, using only a +-- 'Functor f' constraint. +mapGentlyWhenMissing + :: Functor f + => (a -> b) + -> WhenMissing f x a + -> WhenMissing f x b +mapGentlyWhenMissing f t = WhenMissing + { missingSubtree = \m -> fmap f <$> missingSubtree t m + , missingKey = \k x -> fmap f <$> missingKey t k x } +{-# INLINE mapGentlyWhenMissing #-} + + +-- | Map covariantly over a @'WhenMatched' f k x@, using only a +-- 'Functor f' constraint. +mapGentlyWhenMatched + :: Functor f + => (a -> b) + -> WhenMatched f x y a + -> WhenMatched f x y b +mapGentlyWhenMatched f t = + zipWithMaybeAMatched $ \k x y -> fmap f <$> runWhenMatched t k x y +{-# INLINE mapGentlyWhenMatched #-} + + +-- | Map contravariantly over a @'WhenMissing' f _ x@. +-- +-- @since 0.5.9 +lmapWhenMissing :: (b -> a) -> WhenMissing f a x -> WhenMissing f b x +lmapWhenMissing f t = WhenMissing + { missingSubtree = \m -> missingSubtree t (fmap f m) + , missingKey = \k x -> missingKey t k (f x) } +{-# INLINE lmapWhenMissing #-} + + +-- | Map contravariantly over a @'WhenMatched' f _ y z@. +-- +-- @since 0.5.9 +contramapFirstWhenMatched + :: (b -> a) + -> WhenMatched f a y z + -> WhenMatched f b y z +contramapFirstWhenMatched f t = + WhenMatched $ \k x y -> runWhenMatched t k (f x) y +{-# INLINE contramapFirstWhenMatched #-} + + +-- | Map contravariantly over a @'WhenMatched' f x _ z@. +-- +-- @since 0.5.9 +contramapSecondWhenMatched + :: (b -> a) + -> WhenMatched f x a z + -> WhenMatched f x b z +contramapSecondWhenMatched f t = + WhenMatched $ \k x y -> runWhenMatched t k x (f y) +{-# INLINE contramapSecondWhenMatched #-} + + +-- | A tactic for dealing with keys present in one map but not the +-- other in 'merge'. +-- +-- A tactic of type @SimpleWhenMissing x z@ is an abstract +-- representation of a function of type @Key -> x -> Maybe z@. +-- +-- @since 0.5.9 +type SimpleWhenMissing = WhenMissing Identity + + +-- | A tactic for dealing with keys present in both maps in 'merge' +-- or 'mergeA'. +-- +-- A tactic of type @WhenMatched f x y z@ is an abstract representation +-- of a function of type @Key -> x -> y -> f (Maybe z)@. +-- +-- @since 0.5.9 +newtype WhenMatched f x y z = WhenMatched + { matchedKey :: Key -> x -> y -> f (Maybe z) } + + +-- | Along with zipWithMaybeAMatched, witnesses the isomorphism +-- between @WhenMatched f x y z@ and @Key -> x -> y -> f (Maybe z)@. +-- +-- @since 0.5.9 +runWhenMatched :: WhenMatched f x y z -> Key -> x -> y -> f (Maybe z) +runWhenMatched = matchedKey +{-# INLINE runWhenMatched #-} + + +-- | Along with traverseMaybeMissing, witnesses the isomorphism +-- between @WhenMissing f x y@ and @Key -> x -> f (Maybe y)@. +-- +-- @since 0.5.9 +runWhenMissing :: WhenMissing f x y -> Key-> x -> f (Maybe y) +runWhenMissing = missingKey +{-# INLINE runWhenMissing #-} + + +-- | @since 0.5.9 +instance Functor f => Functor (WhenMatched f x y) where + fmap = mapWhenMatched + {-# INLINE fmap #-} + + +-- | @since 0.5.9 +instance Monad f => Category.Category (WhenMatched f x) + where + id = zipWithMatched (\_ _ y -> y) + f . g = + zipWithMaybeAMatched $ \k x y -> do + res <- runWhenMatched g k x y + case res of + Nothing -> pure Nothing + Just r -> runWhenMatched f k x r + {-# INLINE id #-} + {-# INLINE (.) #-} + + +-- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ +-- +-- @since 0.5.9 +instance Monad f => Applicative (WhenMatched f x y) where + pure x = zipWithMatched (\_ _ _ -> x) + fs <*> xs = + zipWithMaybeAMatched $ \k x y -> do + res <- runWhenMatched fs k x y + case res of + Nothing -> pure Nothing + Just r -> (pure $!) . fmap r =<< runWhenMatched xs k x y + {-# INLINE pure #-} + {-# INLINE (<*>) #-} + + +-- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ +-- +-- @since 0.5.9 +instance Monad f => Monad (WhenMatched f x y) where + m >>= f = + zipWithMaybeAMatched $ \k x y -> do + res <- runWhenMatched m k x y + case res of + Nothing -> pure Nothing + Just r -> runWhenMatched (f r) k x y + {-# INLINE (>>=) #-} + + +-- | Map covariantly over a @'WhenMatched' f x y@. +-- +-- @since 0.5.9 +mapWhenMatched + :: Functor f + => (a -> b) + -> WhenMatched f x y a + -> WhenMatched f x y b +mapWhenMatched f (WhenMatched g) = + WhenMatched $ \k x y -> fmap (fmap f) (g k x y) +{-# INLINE mapWhenMatched #-} + + +-- | A tactic for dealing with keys present in both maps in 'merge'. +-- +-- A tactic of type @SimpleWhenMatched x y z@ is an abstract +-- representation of a function of type @Key -> x -> y -> Maybe z@. +-- +-- @since 0.5.9 +type SimpleWhenMatched = WhenMatched Identity + +-- | When a key is found in both maps, drop the key and values. +-- +-- @since 0.8.1 +dropMatched :: Applicative f => WhenMatched f x y z +dropMatched = WhenMatched (\_ _ _ -> pure Nothing) +{-# INLINE dropMatched #-} + +-- | When a key is found in both maps, apply a function to the key +-- and values and use the result in the merged map. +-- +-- > zipWithMatched +-- > :: (Key -> x -> y -> z) +-- > -> SimpleWhenMatched x y z +-- +-- @since 0.5.9 +zipWithMatched + :: Applicative f + => (Key -> x -> y -> z) + -> WhenMatched f x y z +zipWithMatched f = WhenMatched $ \ k x y -> pure . Just $ f k x y +{-# INLINE zipWithMatched #-} + + +-- | When a key is found in both maps, apply a function to the key +-- and values to produce an action and use its result in the merged +-- map. +-- +-- @since 0.5.9 +zipWithAMatched + :: Applicative f + => (Key -> x -> y -> f z) + -> WhenMatched f x y z +zipWithAMatched f = WhenMatched $ \ k x y -> Just <$> f k x y +{-# INLINE zipWithAMatched #-} + + +-- | When a key is found in both maps, apply a function to the key +-- and values and maybe use the result in the merged map. +-- +-- > zipWithMaybeMatched +-- > :: (Key -> x -> y -> Maybe z) +-- > -> SimpleWhenMatched x y z +-- +-- @since 0.5.9 +zipWithMaybeMatched + :: Applicative f + => (Key -> x -> y -> Maybe z) + -> WhenMatched f x y z +zipWithMaybeMatched f = WhenMatched $ \ k x y -> pure $ f k x y +{-# INLINE zipWithMaybeMatched #-} + + +-- | When a key is found in both maps, apply a function to the key +-- and values, perform the resulting action, and maybe use the +-- result in the merged map. +-- +-- This is the fundamental 'WhenMatched' tactic. +-- +-- @since 0.5.9 +zipWithMaybeAMatched + :: (Key -> x -> y -> f (Maybe z)) + -> WhenMatched f x y z +zipWithMaybeAMatched f = WhenMatched $ \ k x y -> f k x y +{-# INLINE zipWithMaybeAMatched #-} + + +-- | Drop all the entries whose keys are missing from the other +-- map. +-- +-- > dropMissing :: SimpleWhenMissing x y +-- +-- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing) +-- +-- but @dropMissing@ is much faster. +-- +-- @since 0.5.9 +dropMissing :: Applicative f => WhenMissing f x y +dropMissing = WhenMissing + { missingSubtree = const (pure Nil) + , missingKey = \_ _ -> pure Nothing } +{-# INLINE dropMissing #-} + + +-- | Preserve, unchanged, the entries whose keys are missing from +-- the other map. +-- +-- > preserveMissing :: SimpleWhenMissing x x +-- +-- prop> preserveMissing = Merge.Lazy.mapMaybeMissing (\_ x -> Just x) +-- +-- but @preserveMissing@ is much faster. +-- +-- @since 0.5.9 +preserveMissing :: Applicative f => WhenMissing f x x +preserveMissing = WhenMissing + { missingSubtree = pure + , missingKey = \_ v -> pure (Just v) } +{-# INLINE preserveMissing #-} + + +-- | Map over the entries whose keys are missing from the other map. +-- +-- > mapMissing :: (k -> x -> y) -> SimpleWhenMissing x y +-- +-- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- +-- but @mapMissing@ is somewhat faster. +-- +-- @since 0.5.9 +mapMissing :: Applicative f => (Key -> x -> y) -> WhenMissing f x y +mapMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapWithKey f m + , missingKey = \k x -> pure $ Just (f k x) } +{-# INLINE mapMissing #-} + + +-- | Map over the entries whose keys are missing from the other +-- map, optionally removing some. This is the most powerful +-- 'SimpleWhenMissing' tactic, but others are usually more efficient. +-- +-- > mapMaybeMissing :: (Key -> x -> Maybe y) -> SimpleWhenMissing x y +-- +-- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) +-- +-- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' +-- operations. +-- +-- @since 0.5.9 +mapMaybeMissing + :: Applicative f => (Key -> x -> Maybe y) -> WhenMissing f x y +mapMaybeMissing f = WhenMissing + { missingSubtree = \m -> pure $! mapMaybeWithKey f m + , missingKey = \k x -> pure $! f k x } +{-# INLINE mapMaybeMissing #-} + + +-- | Filter the entries whose keys are missing from the other map. +-- +-- > filterMissing :: (k -> x -> Bool) -> SimpleWhenMissing x x +-- +-- prop> filterMissing f = Merge.Lazy.mapMaybeMissing $ \k x -> guard (f k x) *> Just x +-- +-- but this should be a little faster. +-- +-- @since 0.5.9 +filterMissing + :: Applicative f => (Key -> x -> Bool) -> WhenMissing f x x +filterMissing f = WhenMissing + { missingSubtree = \m -> pure $! filterWithKey f m + , missingKey = \k x -> pure $! if f k x then Just x else Nothing } +{-# INLINE filterMissing #-} + + +-- | Filter the entries whose keys are missing from the other map +-- using some 'Applicative' action. +-- +-- > filterAMissing f = Merge.Lazy.traverseMaybeMissing $ +-- > \k x -> (\b -> guard b *> Just x) <$> f k x +-- +-- but this should be a little faster. +-- +-- @since 0.5.9 +filterAMissing + :: Applicative f => (Key -> x -> f Bool) -> WhenMissing f x x +filterAMissing f = WhenMissing + { missingSubtree = \m -> filterWithKeyA f m + , missingKey = \k x -> bool Nothing (Just x) <$> f k x } +{-# INLINE filterAMissing #-} + + +-- | \(O(n)\). Filter keys and values using an 'Applicative' predicate. +filterWithKeyA + :: Applicative f => (Key -> a -> f Bool) -> IntMap a -> f (IntMap a) +filterWithKeyA _ Nil = pure Nil +filterWithKeyA f t@(Tip k x) = (\b -> if b then t else Nil) <$> f k x +filterWithKeyA f (Bin p l r) + | signBranch p = liftA2 (flip (bin p)) (filterWithKeyA f r) (filterWithKeyA f l) + | otherwise = liftA2 (bin p) (filterWithKeyA f l) (filterWithKeyA f r) + +-- | This wasn't in Data.Bool until 4.7.0, so we define it here +bool :: a -> a -> Bool -> a +bool f _ False = f +bool _ t True = t + + +-- | Traverse over the entries whose keys are missing from the other +-- map. +-- +-- @since 0.5.9 +traverseMissing + :: Applicative f => (Key -> x -> f y) -> WhenMissing f x y +traverseMissing f = WhenMissing + { missingSubtree = traverseWithKey f + , missingKey = \k x -> Just <$> f k x } +{-# INLINE traverseMissing #-} + + +-- | Traverse over the entries whose keys are missing from the other +-- map, optionally producing values to put in the result. This is +-- the most powerful 'WhenMissing' tactic, but others are usually +-- more efficient. +-- +-- @since 0.5.9 +traverseMaybeMissing + :: Applicative f => (Key -> x -> f (Maybe y)) -> WhenMissing f x y +traverseMaybeMissing f = WhenMissing + { missingSubtree = traverseMaybeWithKey f + , missingKey = f } +{-# INLINE traverseMaybeMissing #-} + + +-- | \(O(n)\). Traverse keys\/values and collect the 'Just' results. +-- +-- @since 0.6.4 +traverseMaybeWithKey + :: Applicative f => (Key -> a -> f (Maybe b)) -> IntMap a -> f (IntMap b) +traverseMaybeWithKey f = go + where + go Nil = pure Nil + go (Tip k x) = maybe Nil (Tip k) <$> f k x + go (Bin p l r) + | signBranch p = liftA2 (flip (bin p)) (go r) (go l) + | otherwise = liftA2 (bin p) (go l) (go r) +{-# INLINE traverseMaybeWithKey #-} + +-- | Merge two maps. +-- +-- 'merge' takes two 'SimpleWhenMissing' tactics, a 'SimpleWhenMatched' tactic +-- and two maps. It uses the tactics to merge the maps. Its behavior +-- is best understood via its fundamental tactics, 'mapMaybeMissing' +-- and 'zipWithMaybeMatched'. +-- +-- Consider +-- +-- @ +-- merge (mapMaybeMissing g1) (mapMaybeMissing g2) (zipWithMaybeMatched f) m1 m2 +-- @ +-- +-- @ +-- g1 k x = if k == 2 then Just ("1" ++ x) else Nothing +-- g2 k x = if k == 3 then Just ("2" ++ x) else Nothing +-- f k x y = if k == 6 then Just ("3" ++ x ++ y) else Nothing +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- m2 = fromList [(3,"g"), (6,"h"), (9,"i"), (12,"j")] +-- @ +-- +-- 'merge' will pass the keys and values to @g1@, @g2@, or @f@ as appropriate, +-- producing a @Maybe@ for each element. +-- +-- @ +-- m1: [ (2, "a"), (4, "b"), (6, "c"), (8, "d"), (10, "e"), (12, "f")] +-- m2: [ (3, "g"), (6, "h"), (9, "i"), (12, "j")] +-- result: [ g1 2 "a", g2 3 "g", g1 4 "b", f 6 "c" "h", g1 8 "d", g2 9 "i", g1 10 "e", f 12 "f" "j"] +-- = [Just "1a", Just "2g", Nothing, Just "3ch", Nothing, Nothing, Nothing, Nothing] +-- @ +-- +-- The result map contains the @Just@ values. +-- +-- >>> merge (mapMaybeMissing g1) (mapMaybeMissing g2) (zipWithMaybeMatched f) m1 m2 +-- fromList [(2,"1a"), (3,"2g"), (6,"3ch")] +-- +-- The other tactics below are optimizations or simplifications of +-- 'mapMaybeMissing' for special cases. Most importantly, +-- +-- * 'dropMissing' drops all the keys. +-- * 'preserveMissing' leaves all the entries alone. +-- +-- When 'merge' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should typically use +-- 'merge' to define your custom combining functions. +-- +-- +-- Examples: +-- +-- prop> unionWithKey f = merge preserveMissing preserveMissing (zipWithMatched f) +-- prop> intersectionWithKey f = merge dropMissing dropMissing (zipWithMatched f) +-- prop> differenceWith f = merge diffPreserve diffDrop f +-- prop> symmetricDifference = merge diffPreserve diffPreserve (\ _ _ _ -> Nothing) +-- prop> mapEachPiece f g h = merge (diffMapWithKey f) (diffMapWithKey g) +-- +-- @since 0.5.9 +merge + :: SimpleWhenMissing a c -- ^ What to do with keys in @m1@ but not @m2@ + -> SimpleWhenMissing b c -- ^ What to do with keys in @m2@ but not @m1@ + -> SimpleWhenMatched a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> IntMap a -- ^ Map @m1@ + -> IntMap b -- ^ Map @m2@ + -> IntMap c +merge g1 g2 f = \m1 m2 -> + runIdentity $ mergeA g1 g2 f m1 m2 +{-# INLINE merge #-} + + +-- | An applicative version of 'merge'. +-- +-- 'mergeA' takes two 'WhenMissing' tactics, a 'WhenMatched' +-- tactic and two maps. It uses the tactics to merge the maps. +-- Its behavior is best understood via its fundamental tactics, +-- 'traverseMaybeMissing' and 'zipWithMaybeAMatched'. +-- +-- Behaves just like 'merge' while allowing @Applicative@ effects. Effects are +-- performed in increasing order of keys. +-- +-- Consider +-- +-- @ +-- mergeA (traverseMaybeMissing g1) +-- (traverseMaybeMissing g2) +-- (zipWithMaybeAMatched f) +-- m1 +-- m2 +-- @ +-- +-- @ +-- g1 k x = let z = if k == 2 then Just ("1" ++ x) else Nothing +-- in z <$ putStrLn ("g1 " ++ show (k, x)) +-- g2 k x = let z = if k == 3 then Just ("2" ++ x) else Nothing +-- in z <$ putStrLn ("g2 " ++ show (k, x)) +-- f k x y = let z = if k == 6 then Just ("3" ++ x ++ y) else Nothing +-- in z <$ putStrLn ("f " ++ show (k, x, y)) +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- m2 = fromList [(3,"g"), (6,"h"), (9,"i"), (12,"j")] +-- @ +-- +-- As with 'merge', the result map is @[(2,"1a"), (3,"2g"), (6,"3ch")]@. +-- Additionally, @g1@, @g2@, and @f@ perform @IO@ effects, printing in +-- increasing order of key. +-- +-- >>> mergeA (traverseMaybeMissing g1) (traverseMaybeMissing g2) (zipWithMaybeAMatched f) m1 m2 +-- g1 (2,"a") +-- g2 (3,"g") +-- g1 (4,"b") +-- f (6,"c","h") +-- g1 (8,"d") +-- g2 (9,"i") +-- g1 (10,"e") +-- f (12,"f","j") +-- fromList [(2,"1a"),(3,"2g"),(6,"3ch")] +-- +-- The other tactics below are optimizations or simplifications of +-- 'traverseMaybeMissing' for special cases. Most importantly, +-- +-- * 'dropMissing' drops all the keys. +-- * 'preserveMissing' leaves all the entries alone. +-- * 'mapMaybeMissing' does not use the 'Applicative' context. +-- +-- When 'mergeA' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should generally only use +-- 'mergeA' to define custom combining functions. +-- +-- === __Examples__ +-- +-- @ +-- data Pair a = Pair !a !a deriving Functor +-- +-- instance Applicative Pair where +-- pure x = Pair x x +-- liftA2 f (Pair x1 y1) (Pair x2 y2) = Pair (f x1 x2) (f y1 y2) +-- +-- -- | Calculate the left-biased union and intersection of two maps. +-- unionIntersection :: IntMap a -> IntMap a -> (IntMap a, IntMap a) +-- unionIntersection m1 m2 = +-- case mergeA preserveAndDropMissing preserveAndDropMissing preserveLeftMatched m1 m2 of +-- Pair mu mi -> (mu, mi) +-- where +-- -- use Pair to build the union and intersection together +-- preserveAndDropMissing = 'whenMissing' (\\_k x -> Pair (Just x) Nothing) (\\m -> Pair m empty) +-- preserveLeftMatched = 'zipWithMaybeAMatched' (\\_k x1 _x2 -> Pair (Just x1) (Just x1)) +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- | Whether the keys of the first map are a subset of the keys of the second map. +-- keysAreSubsetOf :: IntMap a -> IntMap b -> Bool +-- keysAreSubsetOf m1 m2 = +-- getAll (getConst (mergeA isEmpty 'dropMissing' 'dropMatched' m1 m2)) +-- where +-- isEmpty = 'whenMissing' (\\_k _x -> Const (All False)) (\\m -> Const (All (null m))) +-- @ +-- +-- @since 0.5.9 +mergeA + :: (Applicative f) + => WhenMissing f a c -- ^ What to do with keys in @m1@ but not @m2@ + -> WhenMissing f b c -- ^ What to do with keys in @m2@ but not @m1@ + -> WhenMatched f a b c -- ^ What to do with keys in both @m1@ and @m2@ + -> IntMap a -- ^ Map @m1@ + -> IntMap b -- ^ Map @m2@ + -> f (IntMap c) +mergeA + WhenMissing{missingSubtree = g1t, missingKey = g1k} + WhenMissing{missingSubtree = g2t, missingKey = g2k} + WhenMatched{matchedKey = f} + = go + where + go t1 Nil = g1t t1 + go Nil t2 = g2t t2 + + -- This case is already covered below. + -- go (Tip k1 x1) (Tip k2 x2) = mergeTips k1 x1 k2 x2 + + go (Tip k1 x1) t2' = merge2 t2' + where + merge2 t2@(Bin p2 l2 r2) + | nomatch k1 p2 = linkA k1 (subsingletonBy g1k k1 x1) (unPrefix p2) (g2t t2) + | left k1 p2 = binA p2 (merge2 l2) (g2t r2) + | otherwise = binA p2 (g2t l2) (merge2 r2) + merge2 (Tip k2 x2) = mergeTips k1 x1 k2 x2 + merge2 Nil = subsingletonBy g1k k1 x1 + + go t1' (Tip k2 x2) = merge1 t1' + where + merge1 t1@(Bin p1 l1 r1) + | nomatch k2 p1 = linkA (unPrefix p1) (g1t t1) k2 (subsingletonBy g2k k2 x2) + | left k2 p1 = binA p1 (merge1 l1) (g1t r1) + | otherwise = binA p1 (g1t l1) (merge1 r1) + merge1 (Tip k1 x1) = mergeTips k1 x1 k2 x2 + merge1 Nil = subsingletonBy g2k k2 x2 + + go t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> binA p1 (go l1 t2) (g1t r1) + ABR -> binA p1 (g1t l1) (go r1 t2) + BAL -> binA p2 (go t1 l2) (g2t r2) + BAR -> binA p2 (g2t l2) (go t1 r2) + EQL -> binA p1 (go l1 l2) (go r1 r2) + NOM -> linkA (unPrefix p1) (g1t t1) (unPrefix p2) (g2t t2) + + subsingletonBy :: Functor f => (Key -> a -> f (Maybe c)) -> Key -> a -> f (IntMap c) + subsingletonBy gk k x = maybe Nil (Tip k) <$> gk k x + {-# INLINE subsingletonBy #-} + + mergeTips k1 x1 k2 x2 + | k1 == k2 = maybe Nil (Tip k1) <$> f k1 x1 x2 + | k1 < k2 = liftA2 (subdoubleton k1 k2) (g1k k1 x1) (g2k k2 x2) + {- + = link_ k1 k2 <$> subsingletonBy g1k k1 x1 <*> subsingletonBy g2k k2 x2 + -} + | otherwise = liftA2 (subdoubleton k2 k1) (g2k k2 x2) (g1k k1 x1) + {-# INLINE mergeTips #-} + + subdoubleton _ _ Nothing Nothing = Nil + subdoubleton _ k2 Nothing (Just y2) = Tip k2 y2 + subdoubleton k1 _ (Just y1) Nothing = Tip k1 y1 + subdoubleton k1 k2 (Just y1) (Just y2) = link k1 (Tip k1 y1) k2 (Tip k2 y2) + {-# INLINE subdoubleton #-} + + -- | A variant of 'link_' which makes sure to execute side-effects + -- in the right order. + linkA + :: Applicative f + => Int -> f (IntMap a) + -> Int -> f (IntMap a) + -> f (IntMap a) + linkA k1 t1 k2 t2 + | i2w k1 < i2w k2 = binA p t1 t2 + | otherwise = binA p t2 t1 + where + p = branchPrefix k1 k2 + {-# INLINE linkA #-} + + -- A variant of 'bin' that ensures that effects for negative keys are executed + -- first. + binA + :: Applicative f + => Prefix + -> f (IntMap a) + -> f (IntMap a) + -> f (IntMap a) + binA p a b + | signBranch p = liftA2 (flip (bin p)) b a + | otherwise = liftA2 (bin p) a b + {-# INLINE binA #-} +{-# INLINE mergeA #-} + + +{-------------------------------------------------------------------- + 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 l r | signBranch p -> binCheckR p l (go f r) + _ -> go f t + where + go f' (Bin p l r) = binCheckL p (go f' l) r + go f' (Tip k y) = case f' k y of + Just y' -> Tip k y' + Nothing -> Nil + go _ Nil = 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 l r | signBranch p -> binCheckL p (go f l) r + _ -> go f t + where + go f' (Bin p l r) = binCheckR p l (go f' r) + go f' (Tip k y) = case f' k y of + Just y' -> Tip k y' + Nothing -> Nil + go _ Nil = Nil + + +data View a = View {-# UNPACK #-} !Key a !(IntMap a) + +-- | \(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 + _ -> Just $ case maxViewWithKeySure t of + View k v t' -> ((k, v), t') +{-# INLINE maxViewWithKey #-} + +maxViewWithKeySure :: IntMap a -> View a +maxViewWithKeySure t = + case t of + Nil -> error "maxViewWithKeySure Nil" + Bin p l r | signBranch p -> + case go l of View k a l' -> View k a (binCheckL p l' r) + _ -> go t + where + go (Bin p l r) = + case go r of View k a r' -> View k a (binCheckR p l r') + go (Tip k y) = View k y Nil + go Nil = error "maxViewWithKey_go Nil" +-- See note on NOINLINE at minViewWithKeySure +{-# NOINLINE maxViewWithKeySure #-} + +-- | \(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 + _ -> Just $ case minViewWithKeySure t of + View k v t' -> ((k, v), t') +-- We inline this to give GHC the best possible chance of +-- getting rid of the Maybe, pair, and Int constructors, as +-- well as a thunk under the Just. That is, we really want to +-- be certain this inlines! +{-# INLINE minViewWithKey #-} + +minViewWithKeySure :: IntMap a -> View a +minViewWithKeySure t = + case t of + Nil -> error "minViewWithKeySure Nil" + Bin p l r | signBranch p -> + case go r of + View k a r' -> View k a (binCheckR p l r') + _ -> go t + where + go (Bin p l r) = + case go l of View k a l' -> View k a (binCheckL p l' r) + go (Tip k y) = View k y Nil + go Nil = error "minViewWithKey_go Nil" +-- There's never anything significant to be gained by inlining +-- this. Sufficiently recent GHC versions will inline the wrapper +-- anyway, which should be good enough. +{-# NOINLINE minViewWithKeySure #-} + +-- | \(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) + +-- | \(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 = fmap (\((_, x), t') -> (x, t')) (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 = fmap (\((_, x), t') -> (x, t')) (minViewWithKey t) + +-- | \(O(\min(n,W))\). Delete and find the maximal element. +-- +-- Calls 'error' if the map is empty. +-- +-- __Note__: This function is partial. Prefer 'maxViewWithKey'. +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. +-- +-- Calls 'error' if the map is empty. +-- +-- __Note__: This function is partial. Prefer 'minViewWithKey'. +deleteFindMin :: IntMap a -> ((Key, a), IntMap a) +deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey + +-- The KeyValue type is used when returning a key-value pair and helps with +-- GHC optimizations. +-- +-- For lookupMinSure, if the return type is (Int, a), GHC compiles it to a +-- worker $wlookupMinSure :: IntMap a -> (# Int, a #). If the return type is +-- KeyValue a instead, the worker does not box the int and returns +-- (# Int#, a #). +-- For a modern enough GHC (>=9.4), this measure turns out to be unnecessary in +-- this instance. We still use it for older GHCs and to make our intent clear. + +data KeyValue a = KeyValue {-# UNPACK #-} !Key a + +kvToTuple :: KeyValue a -> (Key, a) +kvToTuple (KeyValue k x) = (k, x) +{-# INLINE kvToTuple #-} + +lookupMinSure :: IntMap a -> KeyValue a +lookupMinSure (Tip k v) = KeyValue k v +lookupMinSure (Bin _ l _) = lookupMinSure l +lookupMinSure Nil = error "lookupMinSure Nil" + +-- | \(O(\min(n,W))\). The minimal key of the map. Returns 'Nothing' if the map is empty. +lookupMin :: IntMap a -> Maybe (Key, a) +lookupMin Nil = Nothing +lookupMin (Tip k v) = Just (k,v) +lookupMin (Bin p l r) = + Just $! kvToTuple (lookupMinSure (if signBranch p then r else l)) +{-# INLINE lookupMin #-} -- See Note [Inline lookupMin] in Data.Set.Internal + +-- | \(O(\min(n,W))\). The minimal key of the map. Calls 'error' if the map is empty. +-- +-- __Note__: This function is partial. Prefer 'lookupMin'. +findMin :: IntMap a -> (Key, a) +findMin t + | Just r <- lookupMin t = r + | otherwise = error "findMin: empty map has no minimal element" + +lookupMaxSure :: IntMap a -> KeyValue a +lookupMaxSure (Tip k v) = KeyValue k v +lookupMaxSure (Bin _ _ r) = lookupMaxSure r +lookupMaxSure Nil = error "lookupMaxSure Nil" + +-- | \(O(\min(n,W))\). The maximal key of the map. Returns 'Nothing' if the map is empty. +lookupMax :: IntMap a -> Maybe (Key, a) +lookupMax Nil = Nothing +lookupMax (Tip k v) = Just (k,v) +lookupMax (Bin p l r) = + Just $! kvToTuple (lookupMaxSure (if signBranch p then l else r)) +{-# INLINE lookupMax #-} -- See Note [Inline lookupMin] in Data.Set.Internal + +-- | \(O(\min(n,W))\). The maximal key of the map. Calls 'error' if the map is empty. +-- +-- __Note__: This function is partial. Prefer 'lookupMax'. +findMax :: IntMap a -> (Key, a) +findMax t + | Just r <- lookupMax t = r + | otherwise = error "findMax: empty map has no maximal element" + +-- | \(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(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). + Is this a proper submap? (ie. a submap but not equal). + The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when + @keys m1@ and @keys 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 l1 r1) (Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> GT + ABR -> GT + BAL -> submapCmpLt l2 + BAR -> submapCmpLt r2 + EQL -> submapCmpEq + NOM -> GT -- disjoint + where + submapCmpLt t = case submapCmp predicate t1 t of + GT -> GT + _ -> LT + 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(\min(n, m \log \frac{2^W}{m})), m \leq n\). +-- Is this a submap? +-- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool +isSubmapOf m1 m2 + = isSubmapOfBy (==) m1 m2 + +{- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). + 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 l1 r1) (Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> False + ABR -> False + BAL -> isSubmapOfBy predicate t1 l2 + BAR -> isSubmapOfBy predicate t1 r2 + EQL -> isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2 + NOM -> False +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 = go + where + go (Bin p l r) = Bin p (go l) (go r) + go (Tip k x) = Tip k (f x) + go Nil = Nil + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] map #-} +{-# RULES +"map/map" forall f g xs . map f (map g xs) = map (f . g) xs +"map/coerce" map coerce = coerce + #-} +#endif + +-- | \(O(n)\). Map a function over all values in the map. +-- +-- > let f key x = (show key) ++ ":" ++ x +-- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")] + +mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b +mapWithKey f t + = case t of + Bin p l r -> Bin p (mapWithKey f l) (mapWithKey f r) + Tip k x -> Tip k (f k x) + Nil -> Nil + +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] mapWithKey #-} +{-# RULES +"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = + mapWithKey (\k a -> f k (g k a)) xs +"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = + mapWithKey (\k a -> f k (g a)) xs +"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = + mapWithKey (\k a -> f (g k a)) xs + #-} +#endif + +-- | \(O(n)\). +-- @'traverseWithKey' f 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 l r) + | signBranch p = liftA2 (flip (Bin p)) (go r) (go l) + | otherwise = liftA2 (Bin p) (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 l r + | signBranch p -> + let (a1,r') = mapAccumL f a r + (a2,l') = mapAccumL f a1 l + in (a2,Bin p l' r') + | otherwise -> + let (a1,l') = mapAccumL f a l + (a2,r') = mapAccumL f a1 r + in (a2,Bin p l' r') + Tip k x -> let (a',x') = f a k x in (a',Tip k x') + Nil -> (a,Nil) + +-- | \(O(n)\). The function @'mapAccumRWithKey'@ 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 l r + | signBranch p -> + let (a1,l') = mapAccumRWithKey f a l + (a2,r') = mapAccumRWithKey f a1 r + in (a2,Bin p l' r') + | otherwise -> + let (a1,r') = mapAccumRWithKey f a r + (a2,l') = mapAccumRWithKey f a1 l + in (a2,Bin p 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@. +-- +-- If `f` is monotonically non-decreasing or monotonically non-increasing, this +-- function takes \(O(n)\) time. +-- +-- 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 t = finishB (foldlWithKey' (\b kx x -> insertB (f kx) x b) emptyB t) +{-# INLINABLE mapKeys #-} -- See Note [INLINABLE to expose unfoldings] + +-- | \(O(n \min(n,W))\). +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- If `f` is monotonically non-decreasing or monotonically non-increasing, this +-- function takes \(O(n)\) time. +-- +-- 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" +-- +-- Also see the performance note on 'fromListWith'. + +mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a +mapKeysWith c f t = + finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB t) +{-# INLINABLE mapKeysWith #-} -- See Note [INLINABLE to expose unfoldings] + +-- | \(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@. +-- 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'. +-- +-- __Warning__: This function should be used only if @f@ is monotonically +-- strictly increasing. This precondition is not checked. Use 'mapKeys' if the +-- precondition may not hold. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] + +mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a +mapKeysMonotonic f t = + ascLinkAll (foldlWithKey' (\s kx x -> ascInsert s (f kx) x) MSNada t) +{-# INLINABLE mapKeysMonotonic #-} -- See Note [INLINABLE to expose unfoldings] + +{-------------------------------------------------------------------- + Filter +--------------------------------------------------------------------} +-- | \(O(n)\). Keep 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 = filterWithKey (\_ x -> p x) +{-# INLINE filter #-} + +-- | \(O(n)\). Keep all keys that satisfy some predicate. +-- +-- @ +-- filterKeys p = 'filterWithKey' (\\k _ -> p k) +-- @ +-- +-- > filterKeys (> 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +-- +-- @since 0.8 + +filterKeys :: (Key -> Bool) -> IntMap a -> IntMap a +filterKeys predicate = filterWithKey (\k _ -> predicate k) +{-# INLINE filterKeys #-} + +-- | \(O(n)\). Keep 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 = go + where + go Nil = Nil + go t@(Tip k x) = if predicate k x then t else Nil + go (Bin p l r) = bin p (go l) (go r) +{-# INLINABLE filterWithKey #-} -- See Note [INLINABLE to expose unfoldings] + +-- | \(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 = partitionWithKey (\_ x -> p x) +{-# INLINE partition #-} + +-- | \(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 l r -> + let (l1 :*: l2) = go predicate l + (r1 :*: r2) = go predicate r + in bin p l1 r1 :*: bin p l2 r2 + Tip k x + | predicate k x -> (t :*: Nil) + | otherwise -> (Nil :*: t) + Nil -> (Nil :*: Nil) + +-- | \(O(\min(n,W))\). Take while a predicate on the keys holds. +-- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@. +-- See note at 'spanAntitone'. +-- +-- @ +-- takeWhileAntitone p = 'fromDistinctAscList' . 'Data.List.takeWhile' (p . fst) . 'toList' +-- takeWhileAntitone p = 'filterWithKey' (\\k _ -> p k) +-- @ +-- +-- @since 0.6.7 +takeWhileAntitone :: (Key -> Bool) -> IntMap a -> IntMap a +takeWhileAntitone predicate t = + case t of + Bin p l r + | signBranch p -> + if predicate 0 -- handle negative numbers. + then binCheckL p (go predicate l) r + else go predicate r + _ -> go predicate t + where + go predicate' (Bin p l r) + | predicate' (unPrefix p) = binCheckR p l (go predicate' r) + | otherwise = go predicate' l + go predicate' t'@(Tip ky _) + | predicate' ky = t' + | otherwise = Nil + go _ Nil = Nil + +-- | \(O(\min(n,W))\). Drop while a predicate on the keys holds. +-- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@. +-- See note at 'spanAntitone'. +-- +-- @ +-- dropWhileAntitone p = 'fromDistinctAscList' . 'Data.List.dropWhile' (p . fst) . 'toList' +-- dropWhileAntitone p = 'filterWithKey' (\\k _ -> not (p k)) +-- @ +-- +-- @since 0.6.7 +dropWhileAntitone :: (Key -> Bool) -> IntMap a -> IntMap a +dropWhileAntitone predicate t = + case t of + Bin p l r + | signBranch p -> + if predicate 0 -- handle negative numbers. + then go predicate l + else binCheckR p l (go predicate r) + _ -> go predicate t + where + go predicate' (Bin p l r) + | predicate' (unPrefix p) = go predicate' r + | otherwise = binCheckL p (go predicate' l) r + go predicate' t'@(Tip ky _) + | predicate' ky = Nil + | otherwise = t' + go _ Nil = Nil + +-- | \(O(\min(n,W))\). Divide a map at the point where a predicate on the keys stops holding. +-- The user is responsible for ensuring that for all @Int@s, @j \< k ==\> p j \>= p k@. +-- +-- @ +-- spanAntitone p xs = ('takeWhileAntitone' p xs, 'dropWhileAntitone' p xs) +-- spanAntitone p xs = 'partitionWithKey' (\\k _ -> p k) xs +-- @ +-- +-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the map +-- at some /unspecified/ point. +-- +-- @since 0.6.7 +spanAntitone :: (Key -> Bool) -> IntMap a -> (IntMap a, IntMap a) +spanAntitone predicate t = + case t of + Bin p l r + | signBranch p -> + if predicate 0 -- handle negative numbers. + then + case go predicate l of + (lt :*: gt) -> + let !lt' = binCheckL p lt r + in (lt', gt) + else + case go predicate r of + (lt :*: gt) -> + let !gt' = binCheckR p l gt + in (lt, gt') + _ -> case go predicate t of + (lt :*: gt) -> (lt, gt) + where + go predicate' (Bin p l r) + | predicate' (unPrefix p) + = case go predicate' r of (lt :*: gt) -> binCheckR p l lt :*: gt + | otherwise + = case go predicate' l of (lt :*: gt) -> lt :*: binCheckL p gt r + go predicate' t'@(Tip ky _) + | predicate' ky = (t' :*: Nil) + | otherwise = (Nil :*: t') + go _ 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 l r) + = bin p (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 l r) = + bin p l1 r1 :*: bin p 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 p l r + | signBranch p -> + if k >= 0 -- handle negative numbers. + then + case go k l of + (lt :*: gt) -> + let !lt' = binCheckL p lt r + in (lt', gt) + else + case go k r of + (lt :*: gt) -> + let !gt' = binCheckR p l gt + in (lt, gt') + _ -> case go k t of + (lt :*: gt) -> (lt, gt) + where + go !k' t'@(Bin p l r) + | nomatch k' p = if k' < unPrefix p then Nil :*: t' else t' :*: Nil + | left k' p = case go k' l of (lt :*: gt) -> lt :*: binCheckL p gt r + | otherwise = case go k' r of (lt :*: gt) -> binCheckR p l lt :*: gt + go k' t'@(Tip ky _) + | k' > ky = (t' :*: Nil) + | k' < ky = (Nil :*: t') + | otherwise = (Nil :*: Nil) + go _ Nil = (Nil :*: Nil) + + +type SplitLookup a = StrictTriple (IntMap a) (Maybe a) (IntMap a) + +mapLT :: (IntMap a -> IntMap a) -> SplitLookup a -> SplitLookup a +mapLT f (TripleS lt fnd gt) = TripleS (f lt) fnd gt +{-# INLINE mapLT #-} + +mapGT :: (IntMap a -> IntMap a) -> SplitLookup a -> SplitLookup a +mapGT f (TripleS lt fnd gt) = TripleS lt fnd (f gt) +{-# INLINE mapGT #-} + +-- | \(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 + case t of + Bin p l r + | signBranch p -> + if k >= 0 -- handle negative numbers. + then mapLT (\l' -> binCheckL p l' r) (go k l) + else mapGT (binCheckR p l) (go k r) + _ -> go k t + of TripleS lt fnd gt -> (lt, fnd, gt) + where + go !k' t'@(Bin p l r) + | nomatch k' p = + if k' < unPrefix p + then TripleS Nil Nothing t' + else TripleS t' Nothing Nil + | left k' p = mapGT (\l' -> binCheckL p l' r) (go k' l) + | otherwise = mapLT (binCheckR p l) (go k' r) + go k' t'@(Tip ky y) + | k' > ky = TripleS t' Nothing Nil + | k' < ky = TripleS Nil Nothing t' + | otherwise = TripleS Nil (Just y) Nil + go _ Nil = TripleS 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 + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + go _ Nil = error "foldr.go: Nil" + 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. + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + go !_ Nil = error "foldr'.go: Nil" + 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 + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + go _ Nil = error "foldl.go: Nil" + 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. + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + go !_ Nil = error "foldl'.go: Nil" + go z' (Tip _ x) = f z' x + go z' (Bin _ l r) = go (go z' l) r +{-# INLINE foldl' #-} + +-- See Note [IntMap folds] +foldMap :: Monoid m => (a -> m) -> IntMap a -> m +foldMap f = \t -> -- Use lambda to be inlinable with two arguments. + case t of + Nil -> mempty + Bin p l r +#if MIN_VERSION_base(4,11,0) + | signBranch p -> go r <> go l + | otherwise -> go l <> go r +#else + | signBranch p -> go r `mappend` go l + | otherwise -> go l `mappend` go r +#endif + _ -> go t + where + go Nil = error "foldMap.go: Nil" + go (Tip _ x) = f x +#if MIN_VERSION_base(4,11,0) + go (Bin _ l r) = go l <> go r +#else + go (Bin _ l r) = go l `mappend` go r +#endif +{-# INLINE foldMap #-} + +-- | \(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)" + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + go _ Nil = error "foldrWithKey.go: Nil" + 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. + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t + where + go !_ Nil = error "foldrWithKey'.go: Nil" + 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)" + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + go _ Nil = error "foldlWithKey.go: Nil" + 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. + +-- See Note [IntMap folds] +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 + Nil -> z + Bin p l r + | signBranch p -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t + where + go !_ Nil = error "foldlWithKey'.go: Nil" + 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. +-- +-- @since 0.5.4 + +-- See Note [IntMap folds] +foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m +foldMapWithKey f = \t -> -- Use lambda to be inlinable with two arguments. + case t of + Nil -> mempty + Bin p l r +#if MIN_VERSION_base(4,11,0) + | signBranch p -> go r <> go l + | otherwise -> go l <> go r +#else + | signBranch p -> go r `mappend` go l + | otherwise -> go l `mappend` go r +#endif + _ -> go t + where + go Nil = error "foldMap.go: Nil" + go (Tip kx x) = f kx x +#if MIN_VERSION_base(4,11,0) + go (Bin _ l r) = go l <> go r +#else + go (Bin _ l r) = go l `mappend` go r +#endif +{-# 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)\). 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 +keysSet Nil = IntSet.Nil +keysSet (Tip kx _) = IntSet.singleton kx +keysSet (Bin p l r) + | unPrefix p .&. IntSet.suffixBitMask == 0 + = IntSet.Bin p (keysSet l) (keysSet r) + | otherwise + = IntSet.Tip (unPrefix p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r) + where computeBm !acc (Bin _ l' r') = computeBm (computeBm acc l') r' + computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx + computeBm _ Nil = error "Data.IntSet.keysSet: Nil" + +-- | \(O(n)\). Build a map from an 'IntSet' 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 :: (Key -> a) -> IntSet -> IntMap a +#ifdef __GLASGOW_HASKELL__ +fromSet = + (coerce :: ((Key -> Identity a) -> IntSet -> Identity (IntMap a)) + -> (Key -> a) -> IntSet -> IntMap a) + fromSetA +#else +fromSet f = runIdentity . fromSetA (pure . f) +#endif + +-- | \(O(n)\). Build a map from an 'IntSet' of keys and a function which for +-- each key computes its value, while within an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- > let f k = if k == 0 then Nothing else Just (6 `div` k) +-- > fromSetA f (Data.Set.fromList [1,2,3,4]) == Just (fromList [(1,6),(2,3),(3,2),(4,1)]) +-- > fromSetA f (Data.Set.fromList [0,1,2]) == Nothing +-- +-- @since 0.8.1 +fromSetA :: Applicative f => (Key -> f a) -> IntSet -> f (IntMap a) +fromSetA _ IntSet.Nil = pure Nil +fromSetA f (IntSet.Bin p l r) + | signBranch p = liftA2 (flip (Bin p)) (fromSetA f r) (fromSetA f l) + | otherwise = liftA2 (Bin p) (fromSetA f l) (fromSetA f r) +fromSetA f (IntSet.Tip kx bm) = + treeFromIntSetTip (\kx' -> Tip kx' <$> f kx') (\p -> liftA2 (Bin p)) kx bm +{-# INLINABLE fromSetA #-} + +-- | \(O(n)\). Build a map from an 'IntSet' of keys and a function which for each +-- key optionally computes its value. +-- +-- > let f k = if even k then Just (replicate k 'a') else Nothing +-- > fromSetMaybe f (Data.IntSet.fromList [1,2,3,4]) == fromList [(2,"aa"), (4,"aaaa")] +-- +-- @since 0.8.1 +fromSetMaybe :: (Key -> Maybe a) -> IntSet -> IntMap a +#ifdef __GLASGOW_HASKELL__ +fromSetMaybe = + (coerce :: ((Key -> Identity (Maybe a)) -> IntSet -> Identity (IntMap a)) + -> (Key -> Maybe a) -> IntSet -> IntMap a) + fromSetMaybeA +#else +fromSetMaybe f s = runIdentity (fromSetMaybeA (Identity . f) s) +#endif + +-- | \(O(n)\). Build a map from an 'IntSet' of keys and a function which for +-- each key optionally computes its value in an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- @since 0.8.1 +fromSetMaybeA :: Applicative f => (Key -> f (Maybe a)) -> IntSet -> f (IntMap a) +fromSetMaybeA f = go + where + go IntSet.Nil = pure Nil + go (IntSet.Bin p l r) + | signBranch p = liftA2 (flip (bin p)) (go r) (go l) + | otherwise = liftA2 (bin p) (go l) (go r) + go (IntSet.Tip kx bm) = + treeFromIntSetTip + (\kx' -> maybe Nil (Tip kx') <$> f kx') + (\p -> liftA2 (bin p)) + kx + bm +{-# INLINABLE fromSetMaybeA #-} + +-- Internal helper used by fromSet and friends. +treeFromIntSetTip :: (Int -> a) -> (Prefix -> a -> a -> a) -> Int -> Word -> a +treeFromIntSetTip tipf binf kx bm = buildTree 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 !prefix !bmask bits = case bits of + 0 -> tipf prefix + _ -> case bits `iShiftRL` 1 of + bits2 + | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> + buildTree (prefix + bits2) (bmask `shiftRL` bits2) bits2 + | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 -> + buildTree prefix bmask bits2 + | otherwise -> + binf (Prefix (prefix .|. bits2)) + (buildTree prefix bmask bits2) + (buildTree (prefix + bits2) (bmask `shiftRL` bits2) bits2) +{-# INLINE treeFromIntSetTip #-} + +{-------------------------------------------------------------------- + Lists +--------------------------------------------------------------------} + +#ifdef __GLASGOW_HASKELL__ +-- | @since 0.5.6.2 +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. +-- If the list contains more than one value for the same key, the last value +-- for the key is retained. +-- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- +-- > 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 = finishB (Foldable.foldl' (\b (kx,x) -> insertB kx x b) emptyB xs) +{-# INLINE fromList #-} -- Inline for list fusion + +-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. +-- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"x"), (5,"c")] == fromList [(3, "x"), (5, "cba")] +-- > fromListWith (++) [] == empty +-- +-- Note the reverse ordering of @"cba"@ in the example. +-- +-- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@. +-- +-- See also: 'fromListUpsert' +-- +-- === Performance +-- +-- You should ensure that the given @f@ is fast with this order of arguments. +-- +-- Symmetric functions may be slow in one order, and fast in another. +-- For the common case of collecting values of matching keys in a list, as above: +-- +-- The complexity of @(++) a b@ is \(O(a)\), so it is fast when given a short list as its first argument. +-- Thus: +-- +-- > fromListWith (++) (replicate 1000000 (3, "x")) -- O(n), fast +-- > fromListWith (flip (++)) (replicate 1000000 (3, "x")) -- O(n²), extremely slow +-- +-- because they evaluate as, respectively: +-- +-- > fromList [(3, "x" ++ ("x" ++ "xxxxx..xxxxx"))] -- O(n) +-- > fromList [(3, ("xxxxx..xxxxx" ++ "x") ++ "x")] -- O(n²) +-- +-- Thus, to get good performance with an operation like @(++)@ while also preserving +-- the same order as in the input list, reverse the input: +-- +-- > fromListWith (++) (reverse [(5,"a"), (5,"b"), (5,"c")]) == fromList [(5, "abc")] +-- +-- and it is always fast to combine singleton-list values @[v]@ with @fromListWith (++)@, as in: +-- +-- > fromListWith (++) $ reverse $ map (\(k, v) -> (k, [v])) someListOfTuples + +fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a +fromListWith f xs + = fromListWithKey (\_ x y -> f x y) xs +{-# INLINE fromListWith #-} -- Inline for list fusion + +-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. +-- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- +-- > 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 +-- +-- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromListUpsert' + +fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromListWithKey f xs = + finishB (Foldable.foldl' (\b (kx,x) -> insertWithB (f kx) kx x b) emptyB xs) +{-# INLINE fromListWithKey #-} -- Inline for list fusion + +-- | \(O(n \min(n,W)\). Build a map from a list of key\/value pairs with a +-- combining function. +-- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- +-- The result is equivalent to performing an @upsert@ for every key\/value in +-- the list. +-- +-- @ +-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty' +-- @ +-- +-- > let f x = maybe [x] (x:) +-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")] +-- +-- @since 0.8.1 +fromListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b +fromListUpsert f xs = + finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs) +{-# INLINE fromListUpsert #-} -- INLINE for fusion + +-- | \(O(n)\). Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-decreasing order. This precondition is not checked. Use 'fromList' if the +-- precondition may not hold. +-- +-- > 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 = + ascLinkAll (Foldable.foldl' (\s (ky, y) -> ascInsert s ky y) MSNada xs) +{-# INLINE fromAscList #-} -- Inline for list fusion + +-- | \(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. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-decreasing order. This precondition is not checked. Use 'fromListWith' if +-- the precondition may not hold. +-- +-- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] +-- +-- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' + +fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a +fromAscListWith f xs = fromAscListWithKey (\_ x y -> f x y) xs +{-# INLINE fromAscListWith #-} -- Inline for list fusion + +-- | \(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. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-decreasing order. This precondition is not checked. Use 'fromListWithKey' +-- if the precondition may not hold. +-- +-- > let f key new_value old_value = show key ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromAscListWithKey f [(3,"b"), (3,"a"), (5,"a"), (5,"b"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")] +-- > fromAscListWithKey f [] == empty +-- +-- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' + +-- See Note [fromAscList implementation] +fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a +fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next MSNada xs) + where + next s (!ky, y) = case s of + MSNada -> MSPush ky y Nada + MSPush kx x stk + | kx == ky -> MSPush ky (f ky y x) stk + | otherwise -> let m = branchMask kx ky + in MSPush ky y (ascLinkTop stk kx (Tip kx x) m) +{-# INLINE fromAscListWithKey #-} -- Inline for list fusion + +-- | \(O(n)\). Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-decreasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromAscListUpsert f [(3,'a'), (3,'b'), (5,'c'), (5,'d'), (5,'e')] == fromList [(3,"ba"), (5,"edc")] +-- +-- @since 0.8.1 +fromAscListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b +fromAscListUpsert f xs = ascLinkAll (Foldable.foldl' next MSNada xs) + where + next s (!ky, y) = case s of + MSNada -> MSPush ky (f y Nothing) Nada + MSPush kx x stk + | kx == ky -> MSPush ky (f y (Just x)) stk + | otherwise -> + let m = branchMask kx ky + in MSPush ky (f y Nothing) (ascLinkTop stk kx (Tip kx x) m) +{-# INLINE fromAscListUpsert #-} -- Inline for list fusion + +-- | \(O(n)\). Build a map from a list of key\/value pairs where +-- the keys are in ascending order and all distinct. +-- +-- @fromDistinctAscList = 'fromAscList'@ +-- +-- See warning on 'fromAscList'. +-- +-- This definition exists for backwards compatibility. It offers no advantage +-- over @fromAscList@. +fromDistinctAscList :: [(Key,a)] -> IntMap a +-- Note: There is nothing we can optimize compared to fromAscList. +-- The adjacent key equals check (kx == ky) might seem unnecessary for +-- fromDistinctAscList, but it guards branchMask which has undefined behavior +-- under that case. We could error on kx == ky instead, but that isn't any +-- better. +fromDistinctAscList = fromAscList +{-# INLINE fromDistinctAscList #-} -- Inline for list fusion + +-- | \(O(n)\). Build a map from a list of key\/value pairs where +-- the keys are in descending order. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-increasing order. This precondition is not checked. Use 'fromList' if the +-- precondition may not hold. +-- +-- > fromDescList [(5,"a"), (3,"b")] == fromList [(3,"b"), (5,"a")] +-- > fromDescList [(5,"a"), (5,"b"), (3,"b")] == fromList [(3,"b"), (5,"b")] +-- +-- @since 0.8.1 +fromDescList :: [(Key,a)] -> IntMap a +fromDescList xs = + descLinkAll (Foldable.foldl' (\s (ky, y) -> descInsert ky y s) MSNada xs) +{-# INLINE fromDescList #-} -- Inline for list fusion + +-- | \(O(n)\). Build a map from a descending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-increasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromDescListUpsert f [(5,'a'), (5,'b'), (5,'c'), (3,'d'), (3,'e')] == fromList [(3,"ed"), (5,"cba")] +-- +-- @since 0.8.1 +fromDescListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b +fromDescListUpsert f xs = descLinkAll (Foldable.foldl' next MSNada xs) + where + next s (!ky, y) = case s of + MSNada -> MSPush ky (f y Nothing) Nada + MSPush kx x stk + | kx == ky -> MSPush ky (f y (Just x)) stk + | otherwise -> + let m = branchMask kx ky + in MSPush ky (f y Nothing) (descLinkTop kx (Tip kx x) m stk) +{-# INLINE fromDescListUpsert #-} -- Inline for list fusion + +data Stack a + = Nada + | Push {-# UNPACK #-} !Int !(IntMap a) !(Stack a) + +data MonoState a + = MSNada + | MSPush {-# UNPACK #-} !Key a !(Stack a) + +-- Insert an entry. The key must be >= the last inserted key. If it is equal +-- to the previous key, the previous value is replaced. +ascInsert :: MonoState a -> Int -> a -> MonoState a +ascInsert s !ky y = case s of + MSNada -> MSPush ky y Nada + MSPush kx x stk + | kx == ky -> MSPush ky y stk + | otherwise -> let m = branchMask kx ky + in MSPush ky y (ascLinkTop stk kx (Tip kx x) m) +{-# INLINE ascInsert #-} + +ascLinkTop :: Stack a -> Int -> IntMap a -> Int -> Stack a +ascLinkTop stk !rk r !rm = case stk of + Nada -> Push rm r stk + Push m l stk' + | i2w m < i2w rm -> let p = mask rk m + in ascLinkTop stk' rk (Bin p l r) rm + | otherwise -> Push rm r stk + +ascLinkAll :: MonoState a -> IntMap a +ascLinkAll s = case s of + MSNada -> Nil + MSPush kx x stk -> ascLinkStack stk kx (Tip kx x) +{-# INLINABLE ascLinkAll #-} + +ascLinkStack :: Stack a -> Int -> IntMap a -> IntMap a +ascLinkStack stk !rk r = case stk of + Nada -> r + Push m l stk' + | signBranch p -> Bin p r l + | otherwise -> ascLinkStack stk' rk (Bin p l r) + where + p = mask rk m + +-- Insert an entry. The key must be <= the last inserted key. If it is equal +-- to the previous key, the previous value is replaced. +descInsert :: Int -> a -> MonoState a -> MonoState a +descInsert !ky y s = case s of + MSNada -> MSPush ky y Nada + MSPush kx x stk + | kx == ky -> MSPush ky y stk + | otherwise -> let m = branchMask kx ky + in MSPush ky y (descLinkTop kx (Tip kx x) m stk) +{-# INLINE descInsert #-} + +descLinkTop :: Int -> IntMap a -> Int -> Stack a -> Stack a +descLinkTop !lk l !lm stk = case stk of + Nada -> Push lm l stk + Push m r stk' + | i2w m < i2w lm -> let p = mask lk m + in descLinkTop lk (Bin p l r) lm stk' + | otherwise -> Push lm l stk + +descLinkAll :: MonoState a -> IntMap a +descLinkAll s = case s of + MSNada -> Nil + MSPush kx x stk -> descLinkStack kx (Tip kx x) stk +{-# INLINABLE descLinkAll #-} + +descLinkStack :: Int -> IntMap a -> Stack a -> IntMap a +descLinkStack !lk l stk = case stk of + Nada -> l + Push m r stk' + | signBranch p -> Bin p r l + | otherwise -> descLinkStack lk (Bin p l r) stk' + where + p = mask lk m + +{-------------------------------------------------------------------- + IntMapBuilder +--------------------------------------------------------------------} + +-- Note [IntMapBuilder] +-- ~~~~~~~~~~~~~~~~~~~~ +-- IntMapBuilder serves as an accumulator for element-by-element construction +-- of an IntMap. It can be used in folds to construct IntMaps. This plays nicely +-- with list fusion when the structure folded over is a list, as in fromList and +-- friends. +-- +-- An IntMapBuilder is either empty (BNil) or has the recently inserted Tip +-- together with a stack of trees (BTip). The structure is effectively a +-- [zipper](https://en.wikipedia.org/wiki/Zipper_(data_structure)). It always +-- has its "focus" at the last inserted entry. To insert a new entry, we need +-- to move the focus to the new entry. To do this we move up the stack to the +-- lowest common ancestor of the current position and the position of the +-- new key (implemented as moveUpB), then down to the position of the new key +-- (implemented as moveDownB). +-- +-- When we are done inserting entries, we link the trees up the stack and get +-- the final result. +-- +-- The advantage of this implementation is that we take the shortest path in +-- the tree from one key to the next. Unlike `insert`, we don't need to move +-- up to the root after every insertion. This is very beneficial when we have +-- runs of sorted keys, without many keys already in the tree in that range. +-- If the keys are fully sorted, inserting them all takes O(n) time instead +-- of O(n min(n,W)). But these benefits come at a small cost: when moving up +-- the tree we have to check at every point if it is time to move down. These +-- checks are absent in `insert`. So, in case we need to move up quite a lot, +-- repeated `insert` is slightly faster, but the trade-off is worthwhile since +-- such cases are pathological. + +data IntMapBuilder a + = BNil + | BTip {-# UNPACK #-} !Int a !(BStack a) + +-- BLeft: the IntMap is the left child +-- BRight: the IntMap is the right child +data BStack a + = BNada + | BLeft {-# UNPACK #-} !Prefix !(IntMap a) !(BStack a) + | BRight {-# UNPACK #-} !Prefix !(IntMap a) !(BStack a) + +-- Empty builder. +emptyB :: IntMapBuilder a +emptyB = BNil + +-- Insert a key and value. Replaces the old value if one already exists for +-- the key. +insertB :: Key -> a -> IntMapBuilder a -> IntMapBuilder a +insertB !ky y b = case b of + BNil -> BTip ky y BNada + BTip kx x stk -> case moveToB ky kx x stk of + MoveResult _ stk' -> BTip ky y stk' +{-# INLINE insertB #-} + +-- Insert a key and value. The new value is combined with the old value if one +-- already exists for the key. +insertWithB :: (a -> a -> a) -> Key -> a -> IntMapBuilder a -> IntMapBuilder a +insertWithB f !ky y b = case b of + BNil -> BTip ky y BNada + BTip kx x stk -> case moveToB ky kx x stk of + MoveResult m stk' -> case m of + Nothing -> BTip ky y stk' + Just x' -> BTip ky (f y x') stk' +{-# INLINE insertWithB #-} + +-- Upsert a key-value. The given function is used to generate the value based +-- on the existing value for the key. +upsertB :: (Maybe a -> a) -> Key -> IntMapBuilder a -> IntMapBuilder a +upsertB f !ky b = case b of + BNil -> BTip ky (f Nothing) BNada + BTip kx x stk -> case moveToB ky kx x stk of + MoveResult m stk' -> BTip ky (f m) stk' +{-# INLINE upsertB #-} + +-- GHC >=9.6 supports unpacking sums, so we unpack the Maybe and avoid +-- allocating Justs. GHC optimizes the workers for moveUpB and moveDownB to +-- return (# (# (# #) | a #), BStack a #). +data MoveResult a + = MoveResult +#if __GLASGOW_HASKELL__ >= 906 + {-# UNPACK #-} +#endif + !(Maybe a) + !(BStack a) + +moveToB :: Key -> Key -> a -> BStack a -> MoveResult a +moveToB !ky !kx x !stk + | kx == ky = MoveResult (Just x) stk + | otherwise = moveUpB ky kx (Tip kx x) stk +-- Don't inline this; there is no benefit according to benchmarks. +{-# NOINLINE moveToB #-} + +moveUpB :: Key -> Key -> IntMap a -> BStack a -> MoveResult a +moveUpB !ky !kx !tx stk = case stk of + BNada -> MoveResult Nothing (linkB ky kx tx BNada) + BLeft p l stk' + | nomatch ky p -> moveUpB ky kx (Bin p l tx) stk' + | left ky p -> moveDownB ky l (BRight p tx stk') + | otherwise -> MoveResult Nothing (linkB ky kx tx stk) + BRight p r stk' + | nomatch ky p -> moveUpB ky kx (Bin p tx r) stk' + | left ky p -> MoveResult Nothing (linkB ky kx tx stk) + | otherwise -> moveDownB ky r (BLeft p tx stk') + +moveDownB :: Key -> IntMap a -> BStack a -> MoveResult a +moveDownB !ky tx !stk = case tx of + Bin p l r + | nomatch ky p -> MoveResult Nothing (linkB ky (unPrefix p) tx stk) + | left ky p -> moveDownB ky l (BRight p r stk) + | otherwise -> moveDownB ky r (BLeft p l stk) + Tip kx x + | kx == ky -> MoveResult (Just x) stk + | otherwise -> MoveResult Nothing (linkB ky kx tx stk) + Nil -> error "moveDownB Tip" + +linkB :: Key -> Key -> IntMap a -> BStack a -> BStack a +linkB ky kx tx stk + | i2w ky < i2w kx = BRight p tx stk + | otherwise = BLeft p tx stk + where + p = branchPrefix ky kx +{-# INLINE linkB #-} + +-- Finalize the builder into a Map. +finishB :: IntMapBuilder a -> IntMap a +finishB b = case b of + BNil -> Nil + BTip kx x stk -> finishUpB (Tip kx x) stk +{-# INLINABLE finishB #-} + +finishUpB :: IntMap a -> BStack a -> IntMap a +finishUpB !t stk = case stk of + BNada -> t + BLeft p l stk' -> finishUpB (Bin p l t) stk' + BRight p r stk' -> finishUpB (Bin p t r) stk' + +{-------------------------------------------------------------------- + Eq +--------------------------------------------------------------------} +instance Eq a => Eq (IntMap a) where + (==) = equal + +equal :: Eq a => IntMap a -> IntMap a -> Bool +equal (Bin p1 l1 r1) (Bin p2 l2 r2) + = (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 +{-# INLINABLE equal #-} + +-- | @since 0.5.9 +instance Eq1 IntMap where + liftEq eq = go + where + go (Bin p1 l1 r1) (Bin p2 l2 r2) = p1 == p2 && go l1 l2 && go r1 r2 + go (Tip kx x) (Tip ky y) = kx == ky && eq x y + go Nil Nil = True + go _ _ = False + {-# INLINE liftEq #-} + +{-------------------------------------------------------------------- + Ord +--------------------------------------------------------------------} + +instance Ord a => Ord (IntMap a) where + compare m1 m2 = liftCmp compare m1 m2 + {-# INLINABLE compare #-} + +-- | @since 0.5.9 +instance Ord1 IntMap where + liftCompare = liftCmp + +liftCmp :: (a -> b -> Ordering) -> IntMap a -> IntMap b -> Ordering +liftCmp cmp m1 m2 = case (splitSign m1, splitSign m2) of + ((l1, r1), (l2, r2)) -> case go l1 l2 of + A_LT_B -> LT + A_Prefix_B -> if null r1 then LT else GT + A_EQ_B -> case go r1 r2 of + A_LT_B -> LT + A_Prefix_B -> LT + A_EQ_B -> EQ + B_Prefix_A -> GT + A_GT_B -> GT + B_Prefix_A -> if null r2 then GT else LT + A_GT_B -> GT + where + go t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of + ABL -> case go l1 t2 of + A_Prefix_B -> A_GT_B + A_EQ_B -> B_Prefix_A + o -> o + ABR -> A_LT_B + BAL -> case go t1 l2 of + A_EQ_B -> A_Prefix_B + B_Prefix_A -> A_LT_B + o -> o + BAR -> A_GT_B + EQL -> case go l1 l2 of + A_Prefix_B -> A_GT_B + A_EQ_B -> go r1 r2 + B_Prefix_A -> A_LT_B + o -> o + NOM -> if unPrefix p1 < unPrefix p2 then A_LT_B else A_GT_B + go (Bin _ l1 _) (Tip k2 x2) = case lookupMinSure l1 of + KeyValue k1 x1 -> case compare k1 k2 <> cmp x1 x2 of + LT -> A_LT_B + EQ -> B_Prefix_A + GT -> A_GT_B + go (Tip k1 x1) (Bin _ l2 _) = case lookupMinSure l2 of + KeyValue k2 x2 -> case compare k1 k2 <> cmp x1 x2 of + LT -> A_LT_B + EQ -> A_Prefix_B + GT -> A_GT_B + go (Tip k1 x1) (Tip k2 x2) = case compare k1 k2 <> cmp x1 x2 of + LT -> A_LT_B + EQ -> A_EQ_B + GT -> A_GT_B + go Nil Nil = A_EQ_B + go Nil _ = A_Prefix_B + go _ Nil = B_Prefix_A +{-# INLINE liftCmp #-} + +-- Split into negative and non-negative +splitSign :: IntMap a -> (IntMap a, IntMap a) +splitSign t@(Bin p l r) + | signBranch p = (r, l) + | unPrefix p < 0 = (t, Nil) + | otherwise = (Nil, t) +splitSign t@(Tip k _) + | k < 0 = (t, Nil) + | otherwise = (Nil, t) +splitSign Nil = (Nil, Nil) +{-# INLINE splitSign #-} + +{-------------------------------------------------------------------- + Functor +--------------------------------------------------------------------} + +instance Functor IntMap where + fmap = map + +#ifdef __GLASGOW_HASKELL__ + a <$ Bin p l r = Bin p (a <$ l) (a <$ r) + a <$ Tip k _ = Tip k a + _ <$ Nil = Nil +#endif + +{-------------------------------------------------------------------- + Show +--------------------------------------------------------------------} + +instance Show a => Show (IntMap a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList " . shows (toList m) + +-- | @since 0.5.9 +instance Show1 IntMap where + liftShowsPrec sp sl d m = + showsUnaryWith (liftShowsPrec sp' sl') "fromList" d (toList m) + where + sp' = liftShowsPrec sp sl + sl' = liftShowList sp sl + +{-------------------------------------------------------------------- + Read +--------------------------------------------------------------------} +instance (Read e) => Read (IntMap e) where +#if defined(__GLASGOW_HASKELL__) || defined(__MHS__) + 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 + +-- | @since 0.5.9 +instance Read1 IntMap where + liftReadsPrec rp rl = readsData $ + readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList + where + rp' = liftReadsPrec rp rl + rl' = liftReadList rp rl + +{-------------------------------------------------------------------- + Helpers +--------------------------------------------------------------------} +{-------------------------------------------------------------------- + Link +--------------------------------------------------------------------} + +-- | Link two @IntMap@s. The maps must not be empty. The @Prefix@es of the two +-- maps must be different. @k1@ must share the prefix of @t1@. @p2@ must be the +-- prefix of @t2@. +linkKey :: Key -> IntMap a -> Prefix -> IntMap a -> IntMap a +linkKey k1 t1 p2 t2 = link k1 t1 (unPrefix p2) t2 +{-# INLINE linkKey #-} + +-- | Link two @IntMap@s. The maps must not be empty. The @Prefix@es of the two +-- maps must be different. @k1@ must share the prefix of @t1@ and @k2@ must +-- share the prefix of @t2@. +link :: Int -> IntMap a -> Int -> IntMap a -> IntMap a +link k1 t1 k2 t2 + | i2w k1 < i2w k2 = Bin p t1 t2 + | otherwise = Bin p t2 t1 + where + p = branchPrefix k1 k2 +{-# INLINE link #-} + +{-------------------------------------------------------------------- + @bin@ assures that we never have empty trees within a tree. +--------------------------------------------------------------------} + +bin :: Prefix -> IntMap a -> IntMap a -> IntMap a +bin _ l Nil = l +bin _ Nil r = r +bin p l r = Bin p l r +{-# INLINE bin #-} + +-- binCheckL only checks that the left subtree is non-empty +binCheckL :: Prefix -> IntMap a -> IntMap a -> IntMap a +binCheckL _ Nil r = r +binCheckL p l r = Bin p l r +{-# INLINE binCheckL #-} + +-- binCheckR only checks that the right subtree is non-empty +binCheckR :: Prefix -> IntMap a -> IntMap a -> IntMap a +binCheckR _ l Nil = l +binCheckR p l r = Bin p l r +{-# INLINE binCheckR #-} + +{-------------------------------------------------------------------- + Utilities +--------------------------------------------------------------------} + +-- | \(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 p l r + | signBranch p -> [r, l] + | otherwise -> [l, r] +{-# INLINE splitRoot #-} + + +{-------------------------------------------------------------------- + Debugging +--------------------------------------------------------------------} + +-- | \(O(n \min(n,W))\). 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 \min(n,W))\). 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 l r -> + showsTree wide (withBar rbars) (withEmpty rbars) r . + showWide wide rbars . + showsBars lbars . showString (showBin p) . 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 l r -> + showsBars bars . showString (showBin p) . 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 -> 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 + _ : tl -> showString (concat (reverse tl)) . showString node + +node :: String +node = "+--" + +withBar, withEmpty :: [String] -> [String] +withBar bars = "| ":bars +withEmpty bars = " ":bars + +{-------------------------------------------------------------------- + Notes +--------------------------------------------------------------------} + +-- Note [Okasaki-Gill] +-- ~~~~~~~~~~~~~~~~~~~ +-- +-- The IntMap structure is based on the map described in the paper "Fast +-- Mergeable Integer Maps" by Chris Okasaki and Andy Gill, with some +-- differences. +-- +-- The paper spends most of its time describing a little-endian tree, where the +-- branching is done first on low bits then high bits. It then briefly describes +-- a big-endian tree. The implementation here is big-endian. +-- +-- The definition of Okasaki and Gill's map would be written in Haskell as +-- +-- data Dict a +-- = Empty +-- | Lf !Int a +-- | Br !Int !Int !(Dict a) !(Dict a) +-- +-- Empty is the same as IntMap's Nil, and Lf is the same as Tip. +-- +-- In Br, the first Int is the shared prefix and the second is the mask bit by +-- itself. For the big-endian map, the paper suggests that the prefix be the +-- common prefix, followed by a 0-bit, followed by all 1-bits. This is so that +-- the prefix value can be used as a point of split for binary search. +-- +-- IntMap's Bin corresponds to Br, but is different because it has only one +-- Int (newtyped as Prefix). This describes both prefix and mask, so it is not +-- necessary to store them separately. This value is, in fact, one plus the +-- value suggested for the prefix in the paper. This representation is chosen +-- because it saves one word per Bin without detriment to the efficiency of +-- operations. +-- +-- The implementation of operations such as lookup, insert, union, follow +-- the described implementations on Dict and split into the same cases. For +-- instance, for insert, the three cases on a Br are whether the key belongs +-- outside the map, or it belongs in the left child, or it belongs in the +-- right child. We have the same three cases for a Bin. However, the bitwise +-- operations we use to determine the case is naturally different due to the +-- difference in representation. + +-- Note [IntMap merge complexity] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- The merge algorithm (used for union, intersection, etc.) is adopted from +-- Okasaki-Gill who give the complexity as O(m+n), where m and n are the sizes +-- of the two input maps. This is correct, since we visit all constructors in +-- both maps in the worst case, but we can try to find a tighter bound. +-- +-- Consider that m<=n, i.e. m is the size of the smaller map and n is the size +-- of the larger. It does not matter which map is the first argument. +-- +-- Now we have O(n) as one upper bound for our complexity, since O(n) is the +-- same as O(m+n) for m<=n. +-- +-- Next, consider the smaller map. For this map, we will visit some +-- constructors, plus all the Bins of the larger map that lie in our way. +-- For the former, the worst case is that we visit all constructors, which is +-- O(m). +-- For the latter, the worst case is that we encounter Bins at every point +-- possible. This happens when for every key in the smaller map, the path to +-- that key's Tip in the larger map has a full length of W, with a Bin at every +-- bit position. To maximize the total number of Bins, the paths should be as +-- disjoint as possible. But even if the paths are spread out, at least O(m) +-- Bins are unavoidably shared, which extend up to a depth of lg(m) from the +-- root. Beyond this, the paths may be disjoint. This gives us a total of +-- O(m + m (W - lg m)) = O(m log (2^W / m)). +-- The number of Bins we encounter is also bounded by the total number of Bins, +-- which is n-1, but we already have O(n) as an upper bound. +-- +-- Combining our bounds, we have the final complexity as +-- O(min(n, m log (2^W / m))). +-- +-- Note that +-- * This is similar to the Map merge complexity, which is O(m log (n/m)). +-- * When m is a small constant the term simplifies to O(min(n, W)), which is +-- just the complexity we expect for single operations like insert and delete. + +-- Note [fromAscList implementation] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- fromAscList is an implementation that builds up the result bottom-up +-- in linear time. It maintains a state (MonoState) that gets updated with +-- key-value pairs from the input list one at a time. The state contains the +-- last key-value pair, and a stack of pending trees. +-- +-- For a new key-value pair, the branchMask with the previous key is computed. +-- This represents the depth of the lowest common ancestor that the tree with +-- the previous key, say tl, and the tree with the new key, tr, must have in +-- the final result. Since the keys are in ascending order we expect no more +-- keys in tl, and we can build it by moving up the stack and linking trees. We +-- know when to stop by the branchMask value. We must not link higher than that +-- depth, otherwise instead of tl we will build the parent of tl prematurely +-- before tr is ready. Once the linking is done, tl will be at the top of the +-- stack. +-- +-- We also store the branchMask of a tree with its future right sibling in the +-- stack. This is an optimization, benchmarks show that this is faster than +-- recomputing the branchMask values when linking trees. +-- +-- In the end, we link all the trees remaining in the stack. There is a small +-- catch: negative keys appear in the input before non-negative keys (if they +-- both appear), but the tree with negative keys and the tree with non-negative +-- keys must be the right and left child of the root respectively. So we check +-- for this and link them accordingly. +-- +-- The implementation is defined as a foldl' over the input list, which makes +-- it a good consumer in list fusion. + +-- Note [IntMap folds] +-- ~~~~~~~~~~~~~~~~~~~ +-- Folds on IntMap are defined in a particular way for a few reasons. +-- +-- foldl' :: (a -> b -> a) -> a -> IntMap b -> a +-- foldl' f z = \t -> +-- case t of +-- Nil -> z +-- Bin p l r +-- | signBranch p -> go (go z r) l +-- | otherwise -> go (go z l) r +-- _ -> go z t +-- where +-- go !_ Nil = error "foldl'.go: Nil" +-- go z' (Tip _ x) = f z' x +-- go z' (Bin _ l r) = go (go z' l) r +-- {-# INLINE foldl' #-} +-- +-- 1. We first check if the Bin separates negative and positive keys, and fold +-- over the children accordingly. This check is not inside `go` because it +-- can only happen at the top level and we don't need to check every Bin. +-- 2. We also check for Nil at the top level instead of, say, `go z Nil = z`. +-- That's because `Nil` is also allowed only at the top-level, but more +-- importantly it allows for better optimizations if the `Nil` branch errors +-- in `go`. For example, if we have +-- maximum :: Ord a => IntMap a -> Maybe a +-- maximum = foldl' (\m x -> Just $! maybe x (max x) m) Nothing +-- because `go` certainly returns a `Just` (or errors), CPR analysis will +-- optimize it to return `(# a #)` instead of `Maybe a`. This makes it +-- satisfy the conditions for SpecConstr, which generates two specializations +-- of `go` for `Nothing` and `Just` inputs. Now both `Maybe`s have been +-- optimized out of `go`. +-- 3. The `Tip` is not matched on at the top-level to avoid using `f` more than +-- once. This allows `f` to be inlined into `go` even if `f` is big, since +-- it's likely to be the only place `f` is used, and not inlining `f` means +-- missing out on optimizations. See GHC #25259 for more on this. + +-- Note [INLINABLE to expose unfoldings] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- We have many functions that could be inlined (i.e. not recursive at the top +-- level) but we mark a function with the INLINE pragma only if believe that +-- it will simplify after inlining and improve performance in most situations. +-- Otherwise, inlining just increases code size and compilation times. Fold +-- functions are good examples of functions we surely want to INLINE. +-- +-- For the rest, depending on the function, e.g. if it closes over a +-- user-supplied function, it may improve performance to inline it in certain +-- situations. We want to allow the user to force inlining with GHC.Exts.inline +-- in such situations, so we mark the function as INLINABLE to make its +-- unfolding available in the interface file. +-- +-- For reference see +-- https://downloads.haskell.org/ghc/9.14.1/docs/users_guide/exts/pragmas.html#inlinable-pragma. +-- +-- Note that the user's ability to inline is limited to the body of the +-- function. +-- +-- unionWith f = unionWithKey (\_k x y -> f x y) +-- {-# INLINABLE unionWith #-} +-- unionWithKey f = ...large rhs... +-- {-# INLINABLE unionWithKey #-} +-- +-- Writing `GHC.Exts.inline unionWith` doesn't also inline the body of +-- unionWithKey. If the user wants that, they have to use unionWithKey instead. diff --git a/src/Data/IntMap/Lazy.hs b/src/Data/IntMap/Lazy.hs --- a/src/Data/IntMap/Lazy.hs +++ b/src/Data/IntMap/Lazy.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Lazy @@ -102,19 +100,28 @@ -- * Construction , empty , singleton - , fromSet -- ** From Unordered Lists , fromList , fromListWith , fromListWithKey + , fromListUpsert - -- ** From Ascending Lists + -- ** From Ordered Lists , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList + , fromDescList + , fromDescListUpsert + -- * From @IntSet@ + , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA + -- * Insertion , insert , insertWith @@ -123,10 +130,12 @@ -- * Deletion\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -147,6 +156,7 @@ -- ** Size , IM.null , size + , compareSize -- * Combine @@ -248,12 +258,8 @@ -- * Min\/Max , lookupMin , lookupMax - , findMin - , findMax , deleteMin , deleteMax - , deleteFindMin - , deleteFindMax , updateMin , updateMax , updateMinWithKey @@ -262,6 +268,10 @@ , maxView , minViewWithKey , maxViewWithKey + , findMin + , findMax + , deleteFindMin + , deleteFindMax ) where import Data.IntMap.Internal as IM diff --git a/src/Data/IntMap/Merge/Lazy.hs b/src/Data/IntMap/Merge/Lazy.hs --- a/src/Data/IntMap/Merge/Lazy.hs +++ b/src/Data/IntMap/Merge/Lazy.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Merge.Lazy @@ -42,6 +40,7 @@ , merge -- *** @WhenMatched@ tactics + , dropMatched , zipWithMaybeMatched , zipWithMatched @@ -73,6 +72,7 @@ , traverseMaybeMissing , traverseMissing , filterAMissing + , whenMissing -- *** Covariant maps for tactics , mapWhenMissing diff --git a/src/Data/IntMap/Merge/Strict.hs b/src/Data/IntMap/Merge/Strict.hs --- a/src/Data/IntMap/Merge/Strict.hs +++ b/src/Data/IntMap/Merge/Strict.hs @@ -4,8 +4,6 @@ {-# LANGUAGE Trustworthy #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Merge.Strict @@ -43,6 +41,7 @@ , merge -- *** @WhenMatched@ tactics + , dropMatched , zipWithMaybeMatched , zipWithMatched @@ -74,6 +73,7 @@ , traverseMaybeMissing , traverseMissing , filterAMissing + , Internal.whenMissing -- ** Covariant maps for tactics , mapWhenMissing @@ -96,9 +96,11 @@ , WhenMatched (..) , mergeA , filterAMissing + , dropMatched , runWhenMatched , runWhenMissing ) +import qualified Data.IntMap.Internal as Internal import Data.IntMap.Strict.Internal import Prelude hiding (filter, map, foldl, foldr) diff --git a/src/Data/IntMap/Strict.hs b/src/Data/IntMap/Strict.hs --- a/src/Data/IntMap/Strict.hs +++ b/src/Data/IntMap/Strict.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Trustworthy #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Strict @@ -120,19 +118,28 @@ -- * Construction , empty , singleton - , fromSet -- ** From Unordered Lists , fromList , fromListWith , fromListWithKey + , fromListUpsert - -- ** From Ascending Lists + -- ** From Ordered Lists , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList + , fromDescList + , fromDescListUpsert + -- * From @IntSet@ + , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA + -- * Insertion , insert , insertWith @@ -141,10 +148,12 @@ -- * Deletion\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -165,6 +174,7 @@ -- ** Size , null , size + , compareSize -- * Combine @@ -266,12 +276,8 @@ -- * Min\/Max , lookupMin , lookupMax - , findMin - , findMax , deleteMin , deleteMax - , deleteFindMin - , deleteFindMax , updateMin , updateMax , updateMinWithKey @@ -280,6 +286,10 @@ , maxView , minViewWithKey , maxViewWithKey + , findMin + , findMax + , deleteFindMin + , deleteFindMax ) where import Data.IntMap.Strict.Internal diff --git a/src/Data/IntMap/Strict/Internal.hs b/src/Data/IntMap/Strict/Internal.hs --- a/src/Data/IntMap/Strict/Internal.hs +++ b/src/Data/IntMap/Strict/Internal.hs @@ -1,11 +1,6 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} -{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-} - -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Strict.Internal @@ -64,17 +59,24 @@ , empty , singleton , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA -- ** From Unordered Lists , fromList , fromListWith , fromListWithKey + , fromListUpsert - -- ** From Ascending Lists + -- ** From Ordered Lists , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList + , fromDescList + , fromDescListUpsert -- * Insertion , insert @@ -84,10 +86,12 @@ -- * Deletion\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -108,6 +112,7 @@ -- ** Size , null , size + , compareSize -- * Combine @@ -229,18 +234,31 @@ (lookup,map,filter,foldr,foldl,foldl',null) import Prelude () -import Data.Bits import qualified Data.IntMap.Internal as L import Data.IntSet.Internal.IntTreeCommons - (Key, Prefix(..), nomatch, left, signBranch, mask, branchMask) + (Key, nomatch, left, signBranch, branchMask) import Data.IntMap.Internal ( IntMap (..) , bin - , binCheckLeft - , binCheckRight + , binCheckL + , binCheckR , link , linkKey - , linkWithMask + , MonoState(..) + , Stack(..) + , ascLinkTop + , ascLinkAll + , descInsert + , descLinkTop + , descLinkAll + , IntMapBuilder(..) + , BStack(..) + , emptyB + , insertB + , finishB + , moveToB + , MoveResult(..) + , treeFromIntSetTip , (\\) , (!) @@ -265,6 +283,7 @@ , mergeWithKey' , compose , delete + , pop , deleteMin , deleteMax , deleteFindMax @@ -302,6 +321,7 @@ , spanAntitone , restrictKeys , size + , compareSize , split , splitLookup , splitRoot @@ -313,11 +333,16 @@ , unions , withoutKeys ) +import Data.IntSet.Internal (IntSet) import qualified Data.IntSet.Internal as IntSet -import Utils.Containers.Internal.BitUtil (iShiftRL, shiftLL, shiftRL) -import Utils.Containers.Internal.StrictPair +import Utils.Containers.Internal.Strict (StrictPair(..), toPair) import qualified Data.Foldable as Foldable +import Data.Functor.Identity (Identity (..)) +#ifdef __GLASGOW_HASKELL__ +import Data.Coerce +#endif + {-------------------------------------------------------------------- Construction --------------------------------------------------------------------} @@ -493,8 +518,8 @@ case t of Bin p l r | nomatch k p -> t - | left k p -> binCheckLeft p (updateWithKey f k l) r - | otherwise -> binCheckRight p l (updateWithKey f k r) + | left k p -> binCheckL p (updateWithKey f k l) r + | otherwise -> binCheckR p l (updateWithKey f k r) Tip ky y | k==ky -> case f k y of Just !y' -> Tip ky y' @@ -502,6 +527,26 @@ | otherwise -> t Nil -> Nil +-- | \(O(\min(n,W))\). Update the value at a key or insert a value if the key is +-- not in the map. +-- +-- @ +-- let inc = maybe 1 (+1) +-- upsert inc 100 (fromList [(100,1),(300,2)]) == fromList [(100,2),(300,2)] +-- upsert inc 200 (fromList [(100,1),(300,2)]) == fromList [(100,1),(200,1),(300,2)] +-- @ +-- +-- @since 0.8.1 +upsert :: (Maybe a -> a) -> Key -> IntMap a -> IntMap a +upsert f !k t@(Bin p l r) + | nomatch k p = linkKey k (Tip k $! f Nothing) p t + | left k p = Bin p (upsert f k l) r + | otherwise = Bin p l (upsert f k r) +upsert f !k t@(Tip ky y) + | k == ky = Tip ky $! f (Just y) + | otherwise = link k (Tip k $! f Nothing) ky t +upsert f !k Nil = Tip k $! f Nothing + -- | \(O(\min(n,W))\). Look up and update. -- The function returns original value, if it is updated. -- This is different behavior than 'Data.Map.updateLookupWithKey'. @@ -519,8 +564,8 @@ case t of Bin p l r | nomatch k p -> (Nothing :*: t) - | left k p -> let (found :*: l') = go f k l in (found :*: binCheckLeft p l' r) - | otherwise -> let (found :*: r') = go f k r in (found :*: binCheckRight p l r') + | left k p -> let (found :*: l') = go f k l in (found :*: binCheckL p l' r) + | otherwise -> let (found :*: r') = go f k r in (found :*: binCheckR p l r') Tip ky y | k==ky -> case f k y of Just !y' -> (Just y :*: Tip ky y') @@ -540,8 +585,8 @@ | nomatch k p -> case f Nothing of Nothing -> t Just !x -> linkKey k (Tip k x) p t - | left k p -> binCheckLeft p (alter f k l) r - | otherwise -> binCheckRight p l (alter f k r) + | left k p -> binCheckL p (alter f k l) r + | otherwise -> binCheckR p l (alter f k r) Tip ky y | k==ky -> case f (Just y) of Just !x -> Tip ky x @@ -558,27 +603,26 @@ -- or update a value in an 'IntMap'. In short : @'lookup' k \<$\> 'alterF' f k m = f -- ('lookup' k m)@. -- --- Example: --- --- @ --- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) --- interactiveAlter k m = alterF f k m where --- f Nothing = do --- putStrLn $ show k ++ --- " was not found in the map. Would you like to add it?" --- getUserResponse1 :: IO (Maybe String) --- f (Just old) = do --- putStrLn $ "The key is currently bound to " ++ show old ++ --- ". Would you like to change or delete it?" --- getUserResponse2 :: IO (Maybe String) --- @ --- -- 'alterF' is the most general operation for working with an individual -- key that may or may not be in a given map. - +-- -- Note: 'alterF' is a flipped version of the 'at' combinator from -- 'Control.Lens.At'. -- +-- === Examples +-- +-- @ +-- -- Lookup the value at the key, and also remove the existing value or set a new value. +-- lookupAndSet :: Key -> Maybe a -> IntMap a -> (Maybe a, IntMap a) +-- lookupAndSet k new = alterF (\\old -> (old, new)) k +-- @ +-- +-- @ +-- -- Delete the value at the key. If it is absent the result is Nothing. +-- mustDelete :: Key -> IntMap a -> Maybe (IntMap a) +-- mustDelete = alterF (Nothing <$) +-- @ +-- -- @since 0.5.8 alterF :: Functor f @@ -589,7 +633,7 @@ Nothing -> maybe m (const (delete k m)) mv Just !v' -> insert k v' m where mv = lookup k m - +{-# INLINE alterF #-} {-------------------------------------------------------------------- Union @@ -602,6 +646,7 @@ unionsWith :: Foldable f => (a->a->a) -> f (IntMap a) -> IntMap a unionsWith f ts = Foldable.foldl' (unionWith f) empty ts +{-# INLINE unionsWith #-} -- Inline for list fusion -- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). -- The union with a combining function. @@ -611,8 +656,8 @@ -- Also see the performance note on 'fromListWith'. unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 +unionWith f = unionWithKey (\_ x y -> f x y) +{-# INLINE unionWith #-} -- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). -- The union with a combining function. @@ -624,7 +669,12 @@ 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 + = mergeWithKey' Bin f' id id m1 m2 + where + f' (Tip k1 x1) (Tip _k2 x2) = Tip k1 $! f k1 x1 x2 + f' _ _ = error "not Tip" +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE unionWithKey #-} {-------------------------------------------------------------------- Difference @@ -638,8 +688,8 @@ -- > == 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 +differenceWith f = differenceWithKey (\_ x y -> f x y) +{-# INLINE differenceWith #-} -- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). -- Difference with a combining function. When two equal keys are @@ -654,6 +704,8 @@ differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a differenceWithKey f m1 m2 = mergeWithKey f id (const Nil) m1 m2 +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE differenceWithKey #-} {-------------------------------------------------------------------- Intersection @@ -665,8 +717,8 @@ -- > 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 +intersectionWith f = intersectionWithKey (\_ x y -> f x y) +{-# INLINE intersectionWith #-} -- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). -- The intersection with a combining function. @@ -676,7 +728,12 @@ 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' bin f' (const Nil) (const Nil) m1 m2 + where + f' (Tip k1 x1) (Tip _k2 x2) = Tip k1 $! f k1 x1 x2 + f' _ _ = error "not Tip" +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE intersectionWithKey #-} {-------------------------------------------------------------------- MergeWithKey @@ -722,9 +779,11 @@ 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 + where + combine (Tip k1 x1) (Tip _k2 x2) = case f k1 x1 x2 of + Nothing -> Nil + Just !x -> Tip k1 x + combine _ _ = error "not Tip" {-# INLINE combine #-} {-# INLINE mergeWithKey #-} @@ -739,10 +798,10 @@ updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a updateMinWithKey f t = - case t of Bin p l r | signBranch p -> binCheckRight p l (go f r) + case t of Bin p l r | signBranch p -> binCheckR p l (go f r) _ -> go f t where - go f' (Bin p l r) = binCheckLeft p (go f' l) r + go f' (Bin p l r) = binCheckL p (go f' l) r go f' (Tip k y) = case f' k y of Just !y' -> Tip k y' Nothing -> Nil @@ -755,10 +814,10 @@ updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a updateMaxWithKey f t = - case t of Bin p l r | signBranch p -> binCheckLeft p (go f l) r + case t of Bin p l r | signBranch p -> binCheckL p (go f l) r _ -> go f t where - go f' (Bin p l r) = binCheckRight p l (go f' r) + go f' (Bin p l r) = binCheckR p l (go f' r) go f' (Tip k y) = case f' k y of Just !y' -> Tip k y' Nothing -> Nil @@ -874,6 +933,7 @@ go (Bin p l r) | signBranch p = liftA2 (flip (bin p)) (go r) (go l) | otherwise = liftA2 (bin p) (go l) (go r) +{-# INLINE traverseMaybeWithKey #-} -- | \(O(n)\). The function @'mapAccum'@ threads an accumulating -- argument through the map in ascending order of keys. @@ -937,6 +997,9 @@ -- | \(O(n \min(n,W))\). -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. -- +-- If `f` is monotonically non-decreasing or monotonically non-increasing, this +-- function takes \(O(n)\) time. +-- -- 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@. @@ -947,7 +1010,10 @@ -- Also see the performance note on 'fromListWith'. mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a -mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] +mapKeysWith c f t = + finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB t) +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE mapKeysWith #-} {-------------------------------------------------------------------- Filter @@ -1012,50 +1078,103 @@ Conversions --------------------------------------------------------------------} --- | \(O(n)\). Build a map from a set of keys and a function which for each key --- computes its value. +-- | \(O(n)\). Build a map from an 'IntSet' 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 l r) = Bin p (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 = case bits of - 0 -> Tip prefix $! g prefix - _ -> case bits `iShiftRL` 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 (prefix .|. bits2)) (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2) +fromSet :: (Key -> a) -> IntSet -> IntMap a +#ifdef __GLASGOW_HASKELL__ +fromSet = + (coerce :: ((Key -> Identity a) -> IntSet -> Identity (IntMap a)) + -> (Key -> a) -> IntSet -> IntMap a) + fromSetA +#else +fromSet f = runIdentity . fromSetA (pure . f) +#endif +-- | \(O(n)\). Build a map from an 'IntSet' of keys and a function which for +-- each key computes its value, while within an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- > let f k = if k == 0 then Nothing else Just (6 `div` k) +-- > fromSetA f (Data.Set.fromList [1,2,3,4]) == Just (fromList [(1,6),(2,3),(3,2),(4,1)]) +-- > fromSetA f (Data.Set.fromList [0,1,2]) == Nothing +-- +-- @since 0.8.1 +fromSetA :: Applicative f => (Key -> f a) -> IntSet -> f (IntMap a) +fromSetA _ IntSet.Nil = pure Nil +fromSetA f (IntSet.Bin p l r) + | signBranch p = liftA2 (flip (Bin p)) (fromSetA f r) (fromSetA f l) + | otherwise = liftA2 (Bin p) (fromSetA f l) (fromSetA f r) +fromSetA f (IntSet.Tip kx bm) = + treeFromIntSetTip + (\kx' -> (Tip kx' $!) <$> f kx') + (\p -> liftA2 (Bin p)) + kx + bm +{-# INLINABLE fromSetA #-} + +-- | \(O(n)\). Build a map from an 'IntSet' of keys and a function which for each +-- key optionally computes its value. +-- +-- > let f k = if even k then Just (replicate k 'a') else Nothing +-- > fromSetMaybe f (Data.IntSet.fromList [1,2,3,4]) == fromList [(2,"aa"), (4,"aaaa")] +-- +-- @since 0.8.1 +fromSetMaybe :: (Key -> Maybe a) -> IntSet -> IntMap a +#ifdef __GLASGOW_HASKELL__ +fromSetMaybe = + (coerce :: ((Key -> Identity (Maybe a)) -> IntSet -> Identity (IntMap a)) + -> (Key -> Maybe a) -> IntSet -> IntMap a) + fromSetMaybeA +#else +fromSetMaybe f s = runIdentity (fromSetMaybeA (Identity . f) s) +#endif + +-- | \(O(n)\). Build a map from an 'IntSet' of keys and a function which for +-- each key optionally computes its value in an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- @since 0.8.1 +fromSetMaybeA :: Applicative f => (Key -> f (Maybe a)) -> IntSet -> f (IntMap a) +fromSetMaybeA f = go + where + go IntSet.Nil = pure Nil + go (IntSet.Bin p l r) + | signBranch p = liftA2 (flip (bin p)) (go r) (go l) + | otherwise = liftA2 (bin p) (go l) (go r) + go (IntSet.Tip kx bm) = + treeFromIntSetTip + (\kx' -> maybe Nil (Tip kx' $!) <$> f kx') + (\p -> liftA2 (bin p)) + kx + bm +{-# INLINABLE fromSetMaybeA #-} + {-------------------------------------------------------------------- Lists --------------------------------------------------------------------} -- | \(O(n \min(n,W))\). Create a map from a list of key\/value pairs. -- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- -- > 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 - = Foldable.foldl' ins empty xs - where - ins t (k,x) = insert k x t +fromList xs = finishB (Foldable.foldl' (\b (kx,!x) -> insertB kx x b) emptyB xs) +{-# INLINE fromList #-} -- Inline for list fusion --- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. -- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- -- > fromListWith (++) [(5,"a"), (5,"b"), (3,"x"), (5,"c")] == fromList [(3, "x"), (5, "cba")] -- > fromListWith (++) [] == empty -- @@ -1063,6 +1182,8 @@ -- -- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@. -- +-- See also: 'fromListUpsert' +-- -- === Performance -- -- You should ensure that the given @f@ is fast with this order of arguments. @@ -1093,21 +1214,48 @@ fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a fromListWith f xs = fromListWithKey (\_ x y -> f x y) xs +{-# INLINE fromListWith #-} -- Inline for list fusion --- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. +-- | \(O(n \min(n,W))\). Build a map from a list of key\/value pairs with a combining function. -- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- -- > 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 -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromListUpsert' fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWithKey f xs - = Foldable.foldl' ins empty xs - where - ins t (k,x) = insertWithKey f k x t +fromListWithKey f xs = + finishB (Foldable.foldl' (\b (kx,x) -> insertWithB (f kx) kx x b) emptyB xs) +{-# INLINE fromListWithKey #-} -- Inline for list fusion +-- | \(O(n \min(n,W)\). Build a map from a list of key\/value pairs with a +-- combining function. +-- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. +-- +-- The result is equivalent to performing an @upsert@ for every key\/value in +-- the list. +-- +-- @ +-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty' +-- @ +-- +-- > let f x = maybe [x] (x:) +-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")] +-- +-- @since 0.8.1 +fromListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b +fromListUpsert f xs = + finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs) +{-# INLINE fromListUpsert #-} -- INLINE for fusion + -- | \(O(n)\). Build a map from a list of key\/value pairs where -- the keys are in ascending order. -- @@ -1119,8 +1267,8 @@ -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] fromAscList :: [(Key,a)] -> IntMap a -fromAscList = fromMonoListWithKey Nondistinct (\_ x _ -> x) -{-# NOINLINE fromAscList #-} +fromAscList xs = fromAscListWithKey (\_ x _ -> x) xs +{-# INLINE fromAscList #-} -- Inline for list fusion -- | \(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. @@ -1132,10 +1280,12 @@ -- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f = fromMonoListWithKey Nondistinct (\_ x y -> f x y) -{-# NOINLINE fromAscListWith #-} +fromAscListWith f xs = fromAscListWithKey (\_ x y -> f x y) xs +{-# INLINE fromAscListWith #-} -- Inline for list fusion -- | \(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. @@ -1144,93 +1294,129 @@ -- non-decreasing order. This precondition is not checked. Use 'fromListWithKey' -- if the precondition may not hold. -- --- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] +-- > let f key new_value old_value = show key ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromAscListWithKey f [(3,"b"), (3,"a"), (5,"a"), (5,"b"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")] +-- > fromAscListWithKey f [] == empty -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' +-- See Note [fromAscList implementation] in Data.IntMap.Internal. fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey f = fromMonoListWithKey Nondistinct f -{-# NOINLINE fromAscListWithKey #-} +fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next MSNada xs) + where + next s (!ky, y) = case s of + MSNada -> msPush' ky y Nada + MSPush kx x stk + | kx == ky -> msPush' ky (f ky y x) stk + | otherwise -> let m = branchMask kx ky + in msPush' ky y (ascLinkTop stk kx (Tip kx x) m) + msPush' ky !y = MSPush ky y +{-# INLINE fromAscListWithKey #-} -- Inline for list fusion --- | \(O(n)\). Build a map from a list of key\/value pairs where --- the keys are in ascending order and all distinct. +-- | \(O(n)\). Build a map from an ascending list in linear time with a +-- combining function for equal keys. -- -- __Warning__: This function should be used only if the keys are in --- strictly increasing order. This precondition is not checked. Use 'fromList' +-- non-decreasing order. This precondition is not checked. Use 'fromListUpsert' -- if the precondition may not hold. -- --- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > let f x = maybe [x] (x:) +-- > fromAscListUpsert f [(3,'a'), (3,'b'), (5,'c'), (5,'d'), (5,'e')] == fromList [(3,"ba"), (5,"edc")] +-- +-- @since 0.8.1 +fromAscListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b +fromAscListUpsert f xs = ascLinkAll (Foldable.foldl' next MSNada xs) + where + next s (!ky, y) = case s of + MSNada -> msPush' ky (f y Nothing) Nada + MSPush kx x stk + | kx == ky -> msPush' ky (f y (Just x)) stk + | otherwise -> + let m = branchMask kx ky + in msPush' ky (f y Nothing) (ascLinkTop stk kx (Tip kx x) m) + msPush' ky !y = MSPush ky y +{-# INLINE fromAscListUpsert #-} -- Inline for list fusion +-- | \(O(n)\). Build a map from a list of key\/value pairs where +-- the keys are in ascending order and all distinct. +-- +-- @fromDistinctAscList = 'fromAscList'@ +-- +-- See warning on 'fromAscList'. +-- +-- This definition exists for backwards compatibility. It offers no advantage +-- over @fromAscList@. fromDistinctAscList :: [(Key,a)] -> IntMap a -fromDistinctAscList = fromMonoListWithKey Distinct (\_ x _ -> x) -{-# NOINLINE fromDistinctAscList #-} +-- See Note on Data.IntMap.Internal.fromDistinctAscList. +fromDistinctAscList = fromAscList +{-# INLINE fromDistinctAscList #-} -- Inline for list fusion --- | \(O(n)\). Build a map from a list of key\/value pairs with monotonic keys --- and a combining function. +-- | \(O(n)\). Build a map from a list of key\/value pairs where +-- the keys are in descending order. -- --- The precise conditions under which this function works are subtle: --- For any branch mask, keys with the same prefix w.r.t. the branch --- mask must occur consecutively in the list. +-- __Warning__: This function should be used only if the keys are in +-- non-increasing order. This precondition is not checked. Use 'fromList' if the +-- precondition may not hold. -- --- Also see the performance note on 'fromListWith'. +-- > fromDescList [(5,"a"), (3,"b")] == fromList [(3,"b"), (5,"a")] +-- > fromDescList [(5,"a"), (5,"b"), (3,"b")] == fromList [(3,"b"), (5,"b")] +-- +-- @since 0.8.1 +fromDescList :: [(Key,a)] -> IntMap a +fromDescList xs = + descLinkAll (Foldable.foldl' (\s (!ky, !y) -> descInsert ky y s) MSNada xs) +{-# INLINE fromDescList #-} -- Inline for list fusion -fromMonoListWithKey :: Distinct -> (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromMonoListWithKey distinct f = go +-- | \(O(n)\). Build a map from a descending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-increasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromDescListUpsert f [(5,'a'), (5,'b'), (5,'c'), (3,'d'), (3,'e')] == fromList [(3,"ed"), (5,"cba")] +-- +-- @since 0.8.1 +fromDescListUpsert :: (a -> Maybe b -> b) -> [(Key, a)] -> IntMap b +fromDescListUpsert f xs = descLinkAll (Foldable.foldl' next MSNada xs) where - go [] = Nil - go ((kx,vx) : zs1) = addAll' kx vx zs1 - - -- `addAll'` collects all keys equal to `kx` into a single value, - -- and then proceeds with `addAll`. - -- - -- We want to have the same strictness as fromListWithKey, which is achieved - -- with the bang on vx. - addAll' !kx !vx [] - = Tip kx vx - addAll' !kx !vx ((ky,vy) : zs) - | Nondistinct <- distinct, kx == ky - = addAll' ky (f kx vy vx) zs - -- inlined: | otherwise = addAll kx (Tip kx vx) (ky : zs) - | m <- branchMask kx ky - , Inserted ty zs' <- addMany' m ky vy zs - = addAll kx (linkWithMask m ky ty kx (Tip kx vx)) zs' - - -- for `addAll` and `addMany`, kx is /a/ key inside the tree `tx` - -- `addAll` consumes the rest of the list, adding to the tree `tx` - addAll !_kx !tx [] - = tx - addAll !kx !tx ((ky,vy) : zs) - | m <- branchMask kx ky - , Inserted ty zs' <- addMany' m ky vy zs - = addAll kx (linkWithMask m ky ty kx tx) zs' - - -- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`. - -- - -- We want to have the same strictness as fromListWithKey, which is achieved - -- with the bang on vx. - addMany' !_m !kx !vx [] - = Inserted (Tip kx vx) [] - addMany' !m !kx !vx zs0@((ky,vy) : zs) - | Nondistinct <- distinct, kx == ky - = addMany' m ky (f kx vy vx) zs - -- inlined: | otherwise = addMany m kx (Tip kx vx) (ky : zs) - | mask kx m /= mask ky m - = Inserted (Tip kx vx) zs0 - | mxy <- branchMask kx ky - , Inserted ty zs' <- addMany' mxy ky vy zs - = addMany m kx (linkWithMask mxy ky ty kx (Tip kx vx)) zs' + next s (!ky, y) = case s of + MSNada -> msPush' ky (f y Nothing) Nada + MSPush kx x stk + | kx == ky -> msPush' ky (f y (Just x)) stk + | otherwise -> + let m = branchMask kx ky + in msPush' ky (f y Nothing) (descLinkTop kx (Tip kx x) m stk) + msPush' ky !y = MSPush ky y +{-# INLINE fromDescListUpsert #-} -- Inline for list fusion - -- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `kx`. - addMany !_m !_kx tx [] - = Inserted tx [] - addMany !m !kx tx zs0@((ky,vy) : zs) - | mask kx m /= mask ky m - = Inserted tx zs0 - | mxy <- branchMask kx ky - , Inserted ty zs' <- addMany' mxy ky vy zs - = addMany m kx (linkWithMask mxy ky ty kx tx) zs' -{-# INLINE fromMonoListWithKey #-} +{-------------------------------------------------------------------- + IntMapBuilder +--------------------------------------------------------------------} -data Inserted a = Inserted !(IntMap a) ![(Key,a)] +-- Insert a key and value. The new value is combined with the old value if one +-- already exists for the key. Strict in the inserted value. +insertWithB :: (a -> a -> a) -> Key -> a -> IntMapBuilder a -> IntMapBuilder a +insertWithB f !ky y b = case b of + BNil -> btip' ky y BNada + BTip kx x stk -> case moveToB ky kx x stk of + MoveResult m stk' -> case m of + Nothing -> btip' ky y stk' + Just x' -> btip' ky (f y x') stk' + where + btip' kx !x = BTip kx x +{-# INLINE insertWithB #-} -data Distinct = Distinct | Nondistinct +-- Upsert a key-value. The given function is used to generate the value based +-- on the existing value for the key. +upsertB :: (Maybe a -> a) -> Key -> IntMapBuilder a -> IntMapBuilder a +upsertB f !ky b = case b of + BNil -> btip' ky (f Nothing) BNada + BTip kx x stk -> case moveToB ky kx x stk of + MoveResult m stk' -> btip' ky (f m) stk' + where + btip' kx !x = BTip kx x +{-# INLINE upsertB #-} diff --git a/src/Data/IntSet.hs b/src/Data/IntSet.hs --- a/src/Data/IntSet.hs +++ b/src/Data/IntSet.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntSet @@ -103,12 +101,14 @@ , fromRange , fromAscList , fromDistinctAscList + , fromDescList -- * Insertion , insert -- * Deletion , delete + , pop -- * Generalized insertion/deletion , alterF @@ -122,6 +122,7 @@ , lookupGE , IS.null , size + , compareSize , isSubsetOf , isProperSubsetOf , disjoint @@ -144,6 +145,8 @@ , dropWhileAntitone , spanAntitone + , mapMaybe + , split , splitMember , splitRoot @@ -165,14 +168,14 @@ -- * Min\/Max , lookupMin , lookupMax - , findMin - , findMax , deleteMin , deleteMax - , deleteFindMin - , deleteFindMax , maxView , minView + , findMin + , findMax + , deleteFindMin + , deleteFindMax -- * Conversion diff --git a/src/Data/IntSet/Internal.hs b/src/Data/IntSet/Internal.hs --- a/src/Data/IntSet/Internal.hs +++ b/src/Data/IntSet/Internal.hs @@ -1,9 +1,7 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} #ifdef __GLASGOW_HASKELL__ {-# LANGUAGE DeriveLift #-} -{-# LANGUAGE MagicHash #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE Trustworthy #-} @@ -98,6 +96,7 @@ -- * Query , null , size + , compareSize , member , notMember , lookupLT @@ -114,6 +113,7 @@ , fromRange , insert , delete + , pop , alterF -- * Combine @@ -133,6 +133,8 @@ , dropWhileAntitone , spanAntitone + , mapMaybe + , split , splitMember , splitRoot @@ -175,6 +177,7 @@ , toDescList , fromAscList , fromDistinctAscList + , fromDescList -- * Debugging , showTree @@ -183,6 +186,8 @@ -- * Internals , suffixBitMask , prefixBitMask + , prefixOf + , suffixOf , bitmapOf ) where @@ -198,7 +203,8 @@ import Prelude () import Utils.Containers.Internal.BitUtil (iShiftRL, shiftLL, shiftRL) -import Utils.Containers.Internal.StrictPair +import Utils.Containers.Internal.Strict + (StrictPair(..), StrictTriple(..), toPair) import Data.IntSet.Internal.IntTreeCommons ( Key , Prefix(..) @@ -207,6 +213,7 @@ , signBranch , mask , branchMask + , branchPrefix , TreeTreeBranch(..) , treeTreeBranch , i2w @@ -222,16 +229,16 @@ #if __GLASGOW_HASKELL__ import qualified GHC.Exts -# if !(WORD_SIZE_IN_BITS==64) -import qualified GHC.Int -# endif +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else import Language.Haskell.TH.Syntax (Lift) -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () +# endif #endif import qualified Data.Foldable as Foldable -import Data.Functor.Identity (Identity(..)) infixl 9 \\{-This comment teaches CPP correct behaviour -} @@ -293,14 +300,43 @@ -- -- * In the context of a Tip, the highest (WORD_SIZE - lg(WORD_SIZE)) bits of -- a key are called "prefix" and the lowest lg(WORD_SIZE) bits are called --- "suffix". In Tip kx bm, kx is the shared prefix and bm is a bitmask of the --- suffixes of the keys. In other words, the keys of Tip kx bm are (kx .|. i) --- for every set bit i in bm. --- --- * In Tip kx _, the lowest lg(WORD_SIZE) bits of kx are set to 0. +-- "suffix". In Tip kx bm, kx is the shared prefix of +-- (WORD_SIZE - lg(WORD_SIZE)) bits followed by lg(WORD_SIZE) 0s and bm is a +-- bitmask of the suffixes of the keys. In other words, the keys of Tip kx bm +-- are (kx .|. i) for every set bit i in bm. -- -- * In Tip _ bm, bm is never 0. -- +-- +-- As an example, consider that on a 32-bit system we have a Bin with Prefix +-- +-- 0b00000000000100100000011000000000 +-- ^ mask bit +-- ^^^^^^^^^^^^^^^^^^^^^^ shared prefix +-- +-- The key +-- 0b00000000000100100000010000010100 belongs under this Bin, since it matches +-- the shared prefix. The mask bit is 0, so it belongs in the left child and not +-- the right. +-- +-- The key +-- 0b00000000000100000000010000010100 does not belong under this Bin, since it +-- does not match the shared prefix. +-- +-- +-- Now consider that we have a (Tip kx bm) with kx and bm as +-- +-- 0b00000000001010000000000001000000 0b00000000000100000100000000000100 +-- ^ ^ ^ +-- 20 14 2 -- suffixes +-- ^^^^^^^^^^^^^^^^^^^^^^^^^^^ shared prefix +-- +-- This Tip has 3 keys, which can be recovered by bitwise-or-ing the prefix and +-- the suffixes: +-- +-- 0b00000000001010000000000001000010 +-- 0b00000000001010000000000001001110 +-- 0b00000000001010000000000001010100 #ifdef __GLASGOW_HASKELL__ -- | @since 0.6.6 @@ -311,7 +347,9 @@ instance Monoid IntSet where mempty = empty mconcat = unions +#if !MIN_VERSION_base(4,11,0) mappend = (<>) +#endif -- | @(<>)@ = 'union' -- @@ -355,6 +393,10 @@ {-# INLINE null #-} -- | \(O(n)\). Cardinality of the set. +-- +-- __Note__: Unlike @Data.Set.'Data.Set.size'@, this is /not/ \(O(1)\). +-- +-- See also: 'compareSize' size :: IntSet -> Int size = go 0 where @@ -362,6 +404,26 @@ go acc (Tip _ bm) = acc + popCount bm go acc Nil = acc +-- | \(O(\min(n,c))\). Compare the number of elements in the set to an @Int@. +-- +-- @compareSize m c@ returns the same result as @compare ('size' m) c@ but is +-- more efficient when @c@ is smaller than the size of the set. +-- +-- @since 0.8.1 +compareSize :: IntSet -> Int -> Ordering +compareSize Nil c0 = compare 0 c0 +compareSize _ c0 | c0 <= 0 = GT +compareSize t c0 = compare 0 (go t (c0 - 1)) + where + go (Bin _ _ _) 0 = -1 + go (Bin _ l r) c + | c' < 0 = c' + | otherwise = go r c' + where + c' = go l (c - 1) + go (Tip _ bm) c = c + 1 - popCount bm + go Nil !_ = error "compareSize.go: Nil" + -- | \(O(\min(n,W))\). Is the value a member of the set? -- See Note: Local 'go' functions and capturing. @@ -498,8 +560,7 @@ {-------------------------------------------------------------------- Insert --------------------------------------------------------------------} --- | \(O(\min(n,W))\). Add a value to the set. There is no left- or right bias for --- IntSets. +-- | \(O(\min(n,W))\). Add a value to the set. insert :: Key -> IntSet -> IntSet insert !x = insertBM (prefixOf x) (bitmapOf x) @@ -524,13 +585,46 @@ deleteBM :: Int -> BitMap -> IntSet -> IntSet deleteBM !kx !bm t@(Bin p l r) | nomatch kx p = t - | left kx p = bin p (deleteBM kx bm l) r - | otherwise = bin p l (deleteBM kx bm r) + | left kx p = binCheckL p (deleteBM kx bm l) r + | otherwise = binCheckR p l (deleteBM kx bm r) deleteBM kx bm t@(Tip kx' bm') | kx' == kx = tip kx (bm' .&. complement bm) | otherwise = t deleteBM _ _ Nil = Nil +-- | \(O(\min(n,W))\). Pop an element from the set. +-- +-- Returns @Nothing@ if the element is not a member of the set. Otherwise +-- returns @Just@ the set with the element removed. +-- +-- @ +-- pop 1 (fromList [0,2,4]) == Nothing +-- pop 2 (fromList [0,2,4]) == Just (fromList [0,4]) +-- @ +-- +-- @since 0.8.1 +pop :: Key -> IntSet -> Maybe IntSet +pop x0 t0 = case go x0 t0 of + True :*: t -> Just t + _ -> Nothing + where + -- We use `StrictPair Bool IntSet` instead of a sum to avoid allocations. + -- See Note [Popped impl] in Data.Map.Internal + go !x (Bin p l r) + | nomatch x p = False :*: Nil + | left x p = case go x l of + True :*: l' -> True :*: binCheckL p l' r + q -> q + | otherwise = case go x r of + True :*: r' -> True :*: binCheckR p l r' + q -> q + go !x (Tip ky bmy) + | prefixOf x == ky && bmx .&. bmy /= 0 = True :*: tip ky (bmx `xor` bmy) + | otherwise = False :*: Nil + where + bmx = bitmapOf x + go !_ Nil = False :*: Nil + -- | \(O(\min(n,W))\). @('alterF' f x s)@ can delete or insert @x@ in @s@ depending -- on whether it is already present in @s@. -- @@ -542,30 +636,37 @@ -- -- Note: 'alterF' is a variant of the @at@ combinator from "Control.Lens.At". -- +-- === Examples +-- +-- @ +-- -- Get whether the element is a member, and also insert or remove it. +-- getAndSet :: Key -> Bool -> IntSet -> (Bool, IntSet) +-- getAndSet x new = alterF (\\old -> (old, new)) x +-- @ +-- +-- @ +-- -- Delete the element. If it is absent the result is Nothing. +-- mustDelete :: Key -> IntSet -> Maybe IntSet +-- mustDelete = alterF (\\b -> if b then Just False else Nothing) +-- @ +-- -- @since 0.6.3.1 alterF :: Functor f => (Bool -> f Bool) -> Key -> IntSet -> f IntSet alterF f k s = fmap choose (f member_) where member_ = member k s - - (inserted, deleted) - | member_ = (s , delete k s) - | otherwise = (insert k s, s ) - + inserted = if member_ then s else insert k s + deleted = if member_ then delete k s else s choose True = inserted choose False = deleted -#ifndef __GLASGOW_HASKELL__ -{-# INLINE alterF #-} -#else -{-# INLINABLE [2] alterF #-} +#ifdef __GLASGOW_HASKELL__ +{-# INLINE [2] alterF #-} {-# RULES "alterF/Const" forall k (f :: Bool -> Const a Bool) . alterF f k = \s -> Const . getConst . f $ member k s #-} #endif -{-# SPECIALIZE alterF :: (Bool -> Identity Bool) -> Key -> IntSet -> Identity IntSet #-} - {-------------------------------------------------------------------- Union --------------------------------------------------------------------} @@ -573,7 +674,7 @@ unions :: Foldable f => f IntSet -> IntSet unions xs = Foldable.foldl' union empty xs - +{-# INLINE unions #-} -- Inline for list fusion -- | \(O(\min(n, m \log \frac{2^W}{m})), m \leq n\). -- The union of two sets. @@ -598,8 +699,8 @@ -- Difference between two sets. difference :: IntSet -> IntSet -> IntSet difference t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> bin p1 (difference l1 t2) r1 - ABR -> bin p1 l1 (difference r1 t2) + ABL -> binCheckL p1 (difference l1 t2) r1 + ABR -> binCheckR p1 l1 (difference r1 t2) BAL -> difference t1 l2 BAR -> difference t1 r2 EQL -> bin p1 (difference l1 l2) (difference r1 r2) @@ -709,10 +810,10 @@ symmetricDifference :: IntSet -> IntSet -> IntSet symmetricDifference t1@(Bin p1 l1 r1) t2@(Bin p2 l2 r2) = case treeTreeBranch p1 p2 of - ABL -> bin p1 (symmetricDifference l1 t2) r1 - ABR -> bin p1 l1 (symmetricDifference r1 t2) - BAL -> bin p2 (symmetricDifference t1 l2) r2 - BAR -> bin p2 l2 (symmetricDifference t1 r2) + ABL -> binCheckL p1 (symmetricDifference l1 t2) r1 + ABR -> binCheckR p1 l1 (symmetricDifference r1 t2) + BAL -> binCheckL p2 (symmetricDifference t1 l2) r2 + BAR -> binCheckR p2 l2 (symmetricDifference t1 r2) EQL -> bin p1 (symmetricDifference l1 l2) (symmetricDifference r1 r2) NOM -> link (unPrefix p1) t1 (unPrefix p2) t2 symmetricDifference t1@(Bin _ _ _) t2@(Tip kx2 bm2) = symDiffTip t2 kx2 bm2 t1 @@ -725,8 +826,8 @@ where go t2@(Bin p2 l2 r2) | nomatch kx1 p2 = linkKey kx1 t1 p2 t2 - | left kx1 p2 = bin p2 (go l2) r2 - | otherwise = bin p2 l2 (go r2) + | left kx1 p2 = binCheckL p2 (go l2) r2 + | otherwise = binCheckR p2 l2 (go r2) go t2@(Tip kx2 bm2) | kx1 == kx2 = tip kx1 (bm1 `xor` bm2) | otherwise = link kx1 t1 kx2 t2 @@ -840,7 +941,7 @@ {-------------------------------------------------------------------- Filter --------------------------------------------------------------------} --- | \(O(n)\). Filter all elements that satisfy some predicate. +-- | \(O(n)\). Keep all elements that satisfy some predicate. filter :: (Key -> Bool) -> IntSet -> IntSet filter predicate t = case t of @@ -853,6 +954,20 @@ | otherwise = bm {-# INLINE bitPred #-} +-- | \(O(n \min(n,W))\). Map elements and collect the 'Just' results. +-- +-- If the function is monotonically non-decreasing or monotonically +-- non-increasing, 'mapMaybe' takes \(O(n)\) time. +-- +-- @since 0.8.1 +mapMaybe :: (Key -> Maybe Key) -> IntSet -> IntSet +mapMaybe f t = finishB (foldl' go emptyB t) + where go b x = case f x of + Nothing -> b + Just x' -> insertB x' b +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE mapMaybe #-} + -- | \(O(n)\). partition the set according to some predicate. partition :: (Key -> Bool) -> IntSet -> (IntSet,IntSet) partition predicate0 t0 = toPair $ go predicate0 t0 @@ -887,12 +1002,12 @@ Bin p l r | signBranch p -> if predicate 0 -- handle negative numbers. - then bin p (go predicate l) r + then binCheckL p (go predicate l) r else go predicate r _ -> go predicate t where go predicate' (Bin p l r) - | predicate' (unPrefix p) = bin p l (go predicate' r) + | predicate' (unPrefix p) = binCheckR p l (go predicate' r) | otherwise = go predicate' l go predicate' (Tip kx bm) = tip kx (takeWhileAntitoneBits kx predicate' bm) go _ Nil = Nil @@ -914,12 +1029,12 @@ | signBranch p -> if predicate 0 -- handle negative numbers. then go predicate l - else bin p l (go predicate r) + else binCheckR p l (go predicate r) _ -> go predicate t where go predicate' (Bin p l r) | predicate' (unPrefix p) = go predicate' r - | otherwise = bin p (go predicate' l) r + | otherwise = binCheckL p (go predicate' l) r go predicate' (Tip kx bm) = tip kx (bm `xor` takeWhileAntitoneBits kx predicate' bm) go _ Nil = Nil @@ -944,19 +1059,19 @@ then case go predicate l of (lt :*: gt) -> - let !lt' = bin p lt r + let !lt' = binCheckL p lt r in (lt', gt) else case go predicate r of (lt :*: gt) -> - let !gt' = bin p l gt + let !gt' = binCheckR p l gt in (lt, gt') _ -> case go predicate t of (lt :*: gt) -> (lt, gt) where go predicate' (Bin p l r) - | predicate' (unPrefix p) = case go predicate' r of (lt :*: gt) -> bin p l lt :*: gt - | otherwise = case go predicate' l of (lt :*: gt) -> lt :*: bin p gt r + | predicate' (unPrefix p) = case go predicate' r of (lt :*: gt) -> binCheckR p l lt :*: gt + | otherwise = case go predicate' l of (lt :*: gt) -> lt :*: binCheckL p gt r go predicate' (Tip kx bm) = let bm' = takeWhileAntitoneBits kx predicate' bm in (tip kx bm' :*: tip kx (bm `xor` bm')) go _ Nil = (Nil :*: Nil) @@ -975,20 +1090,20 @@ then case go x l of (lt :*: gt) -> - let !lt' = bin p lt r + let !lt' = binCheckL p lt r in (lt', gt) else case go x r of (lt :*: gt) -> - let !gt' = bin p l gt + let !gt' = binCheckR p l gt in (lt, gt') _ -> case go x t of (lt :*: gt) -> (lt, gt) where go !x' t'@(Bin p l r) | nomatch x' p = if x' < unPrefix p then (Nil :*: t') else (t' :*: Nil) - | left x' p = case go x' l of (lt :*: gt) -> lt :*: bin p gt r - | otherwise = case go x' r of (lt :*: gt) -> bin p l lt :*: gt + | left x' p = case go x' l of (lt :*: gt) -> lt :*: binCheckL p gt r + | otherwise = case go x' r of (lt :*: gt) -> binCheckR p l lt :*: gt go x' t'@(Tip kx' bm) | kx' > x' = (Nil :*: t') -- equivalent to kx' > prefixOf x' @@ -1008,40 +1123,39 @@ if x >= 0 -- handle negative numbers. then case go x l of - (lt, fnd, gt) -> - let !lt' = bin p lt r + TripleS lt fnd gt -> + let !lt' = binCheckL p lt r in (lt', fnd, gt) else case go x r of - (lt, fnd, gt) -> - let !gt' = bin p l gt + TripleS lt fnd gt -> + let !gt' = binCheckR p l gt in (lt, fnd, gt') - _ -> go x t + _ -> case go x t of + TripleS lt fnd gt -> (lt, fnd, gt) where go !x' t'@(Bin p l r) - | nomatch x' p = if x' < unPrefix p then (Nil, False, t') else (t', False, Nil) + | nomatch x' p = if x' < unPrefix p + then TripleS Nil False t' + else TripleS t' False Nil | left x' p = case go x' l of - (lt, fnd, gt) -> - let !gt' = bin p gt r - in (lt, fnd, gt') + TripleS lt fnd gt -> TripleS lt fnd (binCheckL p gt r) | otherwise = case go x' r of - (lt, fnd, gt) -> - let !lt' = bin p l lt - in (lt', fnd, gt) + TripleS lt fnd gt -> TripleS (binCheckR p l lt) fnd gt go x' t'@(Tip kx' bm) - | kx' > x' = (Nil, False, t') + | kx' > x' = TripleS Nil False t' -- equivalent to kx' > prefixOf x' - | kx' < prefixOf x' = (t', False, Nil) + | kx' < prefixOf x' = TripleS t' False Nil | otherwise = let !lt = tip kx' (bm .&. lowerBitmap) !found = (bm .&. bitmapOfx') /= 0 !gt = tip kx' (bm .&. higherBitmap) - in (lt, found, gt) + in TripleS lt found gt where bitmapOfx' = bitmapOf x' lowerBitmap = bitmapOfx' - 1 higherBitmap = complement (lowerBitmap + bitmapOfx') - go _ Nil = (Nil, False, Nil) + go _ Nil = TripleS Nil False Nil {---------------------------------------------------------------------- Min/Max @@ -1052,10 +1166,10 @@ maxView :: IntSet -> Maybe (Key, IntSet) maxView t = case t of Nil -> Nothing - Bin p l r | signBranch p -> case go l of (result, l') -> Just (result, bin p l' r) + Bin p l r | signBranch p -> case go l of (result, l') -> Just (result, binCheckL p l' r) _ -> Just (go t) where - go (Bin p l r) = case go r of (result, r') -> (result, bin p l r') + go (Bin p l r) = case go r of (result, r') -> (result, binCheckR p l r') go (Tip kx bm) = case highestBitSet bm of bi -> (kx + bi, tip kx (bm .&. complement (bitmapOfSuffix bi))) go Nil = error "maxView Nil" @@ -1064,22 +1178,26 @@ minView :: IntSet -> Maybe (Key, IntSet) minView t = case t of Nil -> Nothing - Bin p l r | signBranch p -> case go r of (result, r') -> Just (result, bin p l r') + Bin p l r | signBranch p -> case go r of (result, r') -> Just (result, binCheckR p l r') _ -> Just (go t) where - go (Bin p l r) = case go l of (result, l') -> (result, bin p l' r) + go (Bin p l r) = case go l of (result, l') -> (result, binCheckL p 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) +-- Calls 'error' if the set is empty. +-- +-- __Note__: This function is partial. Prefer 'minView'. 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) +-- Calls 'error' if the set is empty. +-- +-- __Note__: This function is partial. Prefer 'maxView'. deleteFindMax :: IntSet -> (Key, IntSet) deleteFindMax = fromMaybe (error "deleteFindMax: empty set has no maximal element") . maxView @@ -1100,6 +1218,8 @@ -- | \(O(\min(n,W))\). The minimal element of the set. Calls 'error' if the set -- is empty. +-- +-- __Note__: This function is partial. Prefer 'lookupMin'. findMin :: IntSet -> Key findMin t | Just r <- lookupMin t = r @@ -1122,6 +1242,8 @@ -- | \(O(\min(n,W))\). The maximal element of the set. Calls 'error' if the set -- is empty. +-- +-- __Note__: This function is partial. Prefer 'lookupMax'. findMax :: IntSet -> Key findMax t | Just r <- lookupMax t = r @@ -1148,11 +1270,16 @@ -- | \(O(n \min(n,W))\). -- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. -- +-- If `f` is monotonically non-decreasing or monotonically non-increasing, this +-- function takes \(O(n)\) time. +-- -- 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 +map f t = finishB (foldl' (\b x -> insertB (f x) b) emptyB t) +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE map #-} -- | \(O(n)\). The -- @@ -1168,11 +1295,10 @@ -- precondition may not hold. -- -- @since 0.6.3.1 - --- Note that for now the test is insufficient to support any fancier implementation. mapMonotonic :: (Key -> Key) -> IntSet -> IntSet -mapMonotonic f = fromDistinctAscList . List.map f . toAscList - +mapMonotonic f t = ascLinkAll (foldl' (\s x -> ascInsert s (f x)) MSNada t) +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE mapMonotonic #-} {-------------------------------------------------------------------- Fold @@ -1191,13 +1317,18 @@ -- For example, -- -- > toAscList set = foldr (:) [] set + +-- See Note [IntMap folds] in Data.IntMap.Internal 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 p l r | signBranch p -> go (go z l) r -- put negative numbers before - | otherwise -> go (go z r) l - _ -> go z t + case t of + Nil -> z + Bin p l r + | signBranch p -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t where - go z' Nil = z' + go _ Nil = error "foldr.go: Nil" go z' (Tip kx bm) = foldrBits kx f z' bm go z' (Bin _ l r) = go (go z' r) l {-# INLINE foldr #-} @@ -1205,13 +1336,18 @@ -- | \(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. + +-- See Note [IntMap folds] in Data.IntMap.Internal 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 p l r | signBranch p -> go (go z l) r -- put negative numbers before - | otherwise -> go (go z r) l - _ -> go z t + case t of + Nil -> z + Bin p l r + | signBranch p -> go (go z l) r -- put negative numbers before + | otherwise -> go (go z r) l + _ -> go z t where - go !z' Nil = z' + go !_ Nil = error "foldr'.go: Nil" go z' (Tip kx bm) = foldr'Bits kx f z' bm go z' (Bin _ l r) = go (go z' r) l {-# INLINE foldr' #-} @@ -1222,13 +1358,18 @@ -- For example, -- -- > toDescList set = foldl (flip (:)) [] set + +-- See Note [IntMap folds] in Data.IntMap.Internal 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 p l r | signBranch p -> go (go z r) l -- put negative numbers before - | otherwise -> go (go z l) r - _ -> go z t + case t of + Nil -> z + Bin p l r + | signBranch p -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t where - go z' Nil = z' + go _ Nil = error "foldl.go: Nil" go z' (Tip kx bm) = foldlBits kx f z' bm go z' (Bin _ l r) = go (go z' l) r {-# INLINE foldl #-} @@ -1236,13 +1377,18 @@ -- | \(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. + +-- See Note [IntMap folds] in Data.IntMap.Internal 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 p l r | signBranch p -> go (go z r) l -- put negative numbers before - | otherwise -> go (go z l) r - _ -> go z t + case t of + Nil -> z + Bin p l r + | signBranch p -> go (go z r) l -- put negative numbers before + | otherwise -> go (go z l) r + _ -> go z t where - go !z' Nil = z' + go !_ Nil = error "foldl'.go: Nil" go z' (Tip kx bm) = foldl'Bits kx f z' bm go z' (Bin _ l r) = go (go z' l) r {-# INLINE foldl' #-} @@ -1250,6 +1396,8 @@ -- | \(O(n)\). Map the elements in the set to a monoid and combine with @(<>)@. -- -- @since 0.8 + +-- See Note [IntMap folds] in Data.IntMap.Internal foldMap :: Monoid a => (Key -> a) -> IntSet -> a foldMap f = \t -> -- Use lambda t to be inlinable with one argument only. case t of @@ -1261,6 +1409,7 @@ | signBranch p -> go r `mappend` go l -- handle negative numbers | otherwise -> go l `mappend` go r #endif + Nil -> mempty _ -> go t where #if MIN_VERSION_base(4,11,0) @@ -1269,7 +1418,7 @@ go (Bin _ l r) = go l `mappend` go r #endif go (Tip kx bm) = foldMapBits kx f bm - go Nil = mempty + go Nil = error "foldMap.go: Nil" {-# INLINE foldMap #-} {-------------------------------------------------------------------- @@ -1339,11 +1488,12 @@ -- | \(O(n \min(n,W))\). Create a set from a list of integers. +-- +-- If the keys are in sorted order, ascending or descending, this function +-- takes \(O(n)\) time. fromList :: [Key] -> IntSet -fromList xs - = Foldable.foldl' ins empty xs - where - ins t x = insert x t +fromList xs = finishB (Foldable.foldl' (flip insertB) emptyB xs) +{-# INLINE fromList #-} -- Inline for list fusion -- | \(O(n / W)\). Create a set from a range of integers. -- @@ -1355,8 +1505,7 @@ | lx > rx = empty | lp == rp = Tip lp (bitmapOf rx `shiftLL` 1 - bitmapOf lx) | otherwise = - let m = branchMask lx rx - p = Prefix (mask lx m .|. m) + let p = branchPrefix lx rx in if signBranch p -- handle negative numbers then Bin p (goR 0) (goL 0) else Bin p (goL (unPrefix p)) (goR (unPrefix p)) @@ -1403,80 +1552,196 @@ -- __Warning__: This function should be used only if the elements are in -- non-decreasing order. This precondition is not checked. Use 'fromList' if the -- precondition may not hold. + +-- See Note [fromAscList implementation] in Data.IntMap.Internal. fromAscList :: [Key] -> IntSet -fromAscList = fromMonoList -{-# NOINLINE fromAscList #-} +fromAscList xs = ascLinkAll (Foldable.foldl' ascInsert MSNada xs) +{-# INLINE fromAscList #-} -- Inline for list fusion -- | \(O(n)\). Build a set from an ascending list of distinct elements. -- --- __Warning__: This function should be used only if the elements are in --- strictly increasing order. This precondition is not checked. Use 'fromList' --- if the precondition may not hold. +-- @fromDistinctAscList = 'fromAscList'@ +-- +-- See warning on 'fromAscList'. +-- +-- This definition exists for backwards compatibility. It offers no advantage +-- over @fromAscList@. fromDistinctAscList :: [Key] -> IntSet +-- See note on Data.IntMap.Internal.fromDisinctAscList. fromDistinctAscList = fromAscList -{-# INLINE fromDistinctAscList #-} +{-# INLINE fromDistinctAscList #-} -- Inline for list fusion --- | \(O(n)\). Build a set from a monotonic list of elements. +-- | \(O(n)\). Build a set from an descending list of elements. -- --- The precise conditions under which this function works are subtle: --- For any branch mask, keys with the same prefix w.r.t. the branch --- mask must occur consecutively in the list. -fromMonoList :: [Key] -> IntSet -fromMonoList [] = Nil -fromMonoList (kx : zs1) = addAll' (prefixOf kx) (bitmapOf kx) zs1 +-- __Warning__: This function should be used only if the elements are in +-- non-increasing order. This precondition is not checked. Use 'fromList' if the +-- precondition may not hold. +-- +-- @since 0.8.1 + +-- See Note [fromAscList implementation] in Data.IntMap.Internal. +fromDescList :: [Key] -> IntSet +fromDescList xs = descLinkAll (Foldable.foldl' descInsert MSNada xs) +{-# INLINE fromDescList #-} -- Inline for list fusion + +data Stack + = Nada + | Push {-# UNPACK #-} !Int !IntSet !Stack + +data MonoState + = MSNada + | MSPush {-# UNPACK #-} !Int {-# UNPACK #-} !BitMap !Stack + +-- Insert an element. The element must be >= the last inserted element. +ascInsert :: MonoState -> Int -> MonoState +ascInsert s !ky = case s of + MSNada -> MSPush py bmy Nada + MSPush px bmx stk + | px == py -> MSPush py (bmx .|. bmy) stk + | otherwise -> let m = branchMask px py + in MSPush py bmy (ascLinkTop stk px (Tip px bmx) m) where - -- `addAll'` collects all keys with the prefix `px` into a single - -- bitmap, and then proceeds with `addAll`. - addAll' !px !bm [] - = Tip px bm - addAll' !px !bm (ky : zs) - | px == prefixOf ky - = addAll' px (bm .|. bitmapOf ky) zs - -- inlined: | otherwise = addAll px (Tip px bm) (ky : zs) - | py <- prefixOf ky - , m <- branchMask px py - , Inserted ty zs' <- addMany' m py (bitmapOf ky) zs - = addAll px (linkWithMask m py ty px (Tip px bm)) zs' + py = prefixOf ky + bmy = bitmapOf ky +{-# INLINE ascInsert #-} - -- for `addAll` and `addMany`, px is /a/ prefix inside the tree `tx` - -- `addAll` consumes the rest of the list, adding to the tree `tx` - addAll !_px !tx [] - = tx - addAll !px !tx (ky : zs) - | py <- prefixOf ky - , m <- branchMask px py - , Inserted ty zs' <- addMany' m py (bitmapOf ky) zs - = addAll px (linkWithMask m py ty px tx) zs' +ascLinkTop :: Stack -> Int -> IntSet -> Int -> Stack +ascLinkTop stk !rk r !rm = case stk of + Nada -> Push rm r stk + Push m l stk' + | i2w m < i2w rm -> let p = mask rk m + in ascLinkTop stk' rk (Bin p l r) rm + | otherwise -> Push rm r stk - -- `addMany'` is similar to `addAll'`, but proceeds with `addMany'`. - addMany' !_m !px !bm [] - = Inserted (Tip px bm) [] - addMany' !m !px !bm zs0@(ky : zs) - | px == prefixOf ky - = addMany' m px (bm .|. bitmapOf ky) zs - -- inlined: | otherwise = addMany m px (Tip px bm) (ky : zs) - | mask px m /= mask ky m - = Inserted (Tip (prefixOf px) bm) zs0 - | py <- prefixOf ky - , mxy <- branchMask px py - , Inserted ty zs' <- addMany' mxy py (bitmapOf ky) zs - = addMany m px (linkWithMask mxy py ty px (Tip px bm)) zs' +ascLinkAll :: MonoState -> IntSet +ascLinkAll s = case s of + MSNada -> Nil + MSPush px bmx stk -> ascLinkStack stk px (Tip px bmx) +{-# INLINABLE ascLinkAll #-} - -- `addAll` adds to `tx` all keys whose prefix w.r.t. `m` agrees with `px`. - addMany !_m !_px tx [] - = Inserted tx [] - addMany !m !px tx zs0@(ky : zs) - | mask px m /= mask ky m - = Inserted tx zs0 - | py <- prefixOf ky - , mxy <- branchMask px py - , Inserted ty zs' <- addMany' mxy py (bitmapOf ky) zs - = addMany m px (linkWithMask mxy py ty px tx) zs' -{-# INLINE fromMonoList #-} +ascLinkStack :: Stack -> Int -> IntSet -> IntSet +ascLinkStack stk !rk r = case stk of + Nada -> r + Push m l stk' + | signBranch p -> Bin p r l + | otherwise -> ascLinkStack stk' rk (Bin p l r) + where + p = mask rk m -data Inserted = Inserted !IntSet ![Key] +-- Insert an element. The element must be <= the last inserted element. +descInsert :: MonoState -> Int -> MonoState +descInsert s !ky = case s of + MSNada -> MSPush py bmy Nada + MSPush px bmx stk + | px == py -> MSPush py (bmx .|. bmy) stk + | otherwise -> let m = branchMask px py + in MSPush py bmy (descLinkTop px (Tip px bmx) m stk) + where + py = prefixOf ky + bmy = bitmapOf ky +{-# INLINE descInsert #-} +descLinkTop :: Int -> IntSet -> Int -> Stack -> Stack +descLinkTop !lk l !lm stk = case stk of + Nada -> Push lm l stk + Push m r stk' + | i2w m < i2w lm -> let p = mask lk m + in descLinkTop lk (Bin p l r) lm stk' + | otherwise -> Push lm l stk + +descLinkAll :: MonoState -> IntSet +descLinkAll s = case s of + MSNada -> Nil + MSPush px bmx stk -> descLinkStack px (Tip px bmx) stk +{-# INLINABLE descLinkAll #-} + +descLinkStack :: Int -> IntSet -> Stack -> IntSet +descLinkStack !lk l stk = case stk of + Nada -> l + Push m r stk' + | signBranch p -> Bin p r l + | otherwise -> descLinkStack lk (Bin p l r) stk' + where + p = mask lk m + {-------------------------------------------------------------------- + IntSetBuilder +--------------------------------------------------------------------} + +-- See Note [IntMapBuilder] in Data.IntMap.Internal. + +data IntSetBuilder + = BNil + | BTip {-# UNPACK #-} !Int {-# UNPACK #-} !BitMap !BStack + +-- BLeft: the IntMap is the left child +-- BRight: the IntMap is the right child +data BStack + = BNada + | BLeft {-# UNPACK #-} !Prefix !IntSet !BStack + | BRight {-# UNPACK #-} !Prefix !IntSet !BStack + +-- Empty builder. +emptyB :: IntSetBuilder +emptyB = BNil + +-- Insert an element. +insertB :: Key -> IntSetBuilder -> IntSetBuilder +insertB !ky b = case b of + BNil -> BTip py bmy BNada + BTip px bmx stk + | px == py -> BTip py (bmx .|. bmy) stk + | otherwise -> insertUpB py bmy px (Tip px bmx) stk + where + py = prefixOf ky + bmy = bitmapOf ky +{-# INLINE insertB #-} + +insertUpB :: Int -> BitMap -> Int -> IntSet -> BStack -> IntSetBuilder +insertUpB !py !bmy !px !tx stk = case stk of + BNada -> BTip py bmy (linkB py px tx BNada) + BLeft p l stk' + | nomatch py p -> insertUpB py bmy px (Bin p l tx) stk' + | left py p -> insertDownB py bmy l (BRight p tx stk') + | otherwise -> BTip py bmy (linkB py px tx stk) + BRight p r stk' + | nomatch py p -> insertUpB py bmy px (Bin p tx r) stk' + | left py p -> BTip py bmy (linkB py px tx stk) + | otherwise -> insertDownB py bmy r (BLeft p tx stk') + +insertDownB :: Int -> BitMap -> IntSet -> BStack -> IntSetBuilder +insertDownB !py !bmy tx !stk = case tx of + Bin p l r + | nomatch py p -> BTip py bmy (linkB py (unPrefix p) tx stk) + | left py p -> insertDownB py bmy l (BRight p r stk) + | otherwise -> insertDownB py bmy r (BLeft p l stk) + Tip px bmx + | px == py -> BTip py (bmx .|. bmy) stk + | otherwise -> BTip py bmy (linkB py px tx stk) + Nil -> error "insertDownB Tip" + +linkB :: Key -> Key -> IntSet -> BStack -> BStack +linkB ky kx tx stk + | i2w ky < i2w kx = BRight p tx stk + | otherwise = BLeft p tx stk + where + p = branchPrefix ky kx +{-# INLINE linkB #-} + +-- Finalize the builder into an IntSet. +finishB :: IntSetBuilder -> IntSet +finishB b = case b of + BNil -> Nil + BTip px bmx stk -> finishUpB (Tip px bmx) stk +{-# INLINABLE finishB #-} + +finishUpB :: IntSet -> BStack -> IntSet +finishUpB !t stk = case stk of + BNada -> t + BLeft p l stk' -> finishUpB (Bin p l t) stk' + BRight p r stk' -> finishUpB (Bin p t r) stk' + +{-------------------------------------------------------------------- Eq --------------------------------------------------------------------} instance Eq IntSet where @@ -1711,17 +1976,12 @@ -- sets must be different. @k1@ must share the prefix of @t1@ and @k2@ must -- share the prefix of @t2@. link :: Int -> IntSet -> Int -> IntSet -> IntSet -link k1 t1 k2 t2 = linkWithMask (branchMask k1 k2) k1 t1 k2 t2 -{-# INLINE link #-} - --- `linkWithMask` is useful when the `branchMask` has already been computed -linkWithMask :: Int -> Key -> IntSet -> Key -> IntSet -> IntSet -linkWithMask m k1 t1 k2 t2 +link k1 t1 k2 t2 | i2w k1 < i2w k2 = Bin p t1 t2 | otherwise = Bin p t2 t1 where - p = Prefix (mask k1 m .|. m) -{-# INLINE linkWithMask #-} + p = branchPrefix k1 k2 +{-# INLINE link #-} {-------------------------------------------------------------------- @bin@ assures that we never have empty trees within a tree. @@ -1731,6 +1991,18 @@ bin _ Nil r = r bin p l r = Bin p l r {-# INLINE bin #-} + +-- binCheckL only checks that the left subtree is non-empty +binCheckL :: Prefix -> IntSet -> IntSet -> IntSet +binCheckL _ Nil r = r +binCheckL p l r = Bin p l r +{-# INLINE binCheckL #-} + +-- binCheckR only checks that the right subtree is non-empty +binCheckR :: Prefix -> IntSet -> IntSet -> IntSet +binCheckR _ l Nil = l +binCheckR p l r = Bin p l r +{-# INLINE binCheckR #-} {-------------------------------------------------------------------- @tip@ assures that we never have empty bitmaps within a tree. diff --git a/src/Data/IntSet/Internal/IntTreeCommons.hs b/src/Data/IntSet/Internal/IntTreeCommons.hs --- a/src/Data/IntSet/Internal/IntTreeCommons.hs +++ b/src/Data/IntSet/Internal/IntTreeCommons.hs @@ -35,17 +35,22 @@ , treeTreeBranch , mask , branchMask + , branchPrefix , i2w , Order(..) ) where import Data.Bits (Bits(..), countLeadingZeros) -import Utils.Containers.Internal.BitUtil (wordSize) +import Utils.Containers.Internal.BitUtil (iShiftRL) #ifdef __GLASGOW_HASKELL__ +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else import Language.Haskell.TH.Syntax (Lift) -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () +# endif #endif @@ -144,19 +149,26 @@ signBranch p = unPrefix p == (minBound :: Int) {-# INLINE signBranch #-} --- | The prefix of key @i@ up to (but not including) the switching --- bit @m@. -mask :: Key -> Int -> Int -mask i m = i .&. ((-m) `xor` m) +-- | The prefix of @Int@ @i@ up to the switching bit @m@. +mask :: Int -> Int -> Prefix +mask i m = Prefix ((i .&. negate m) .|. m) {-# INLINE mask #-} --- | The first switching bit where the two prefixes disagree. +-- | The first switching bit where the two @Int@s disagree. -- --- Precondition for defined behavior: p1 /= p2 +-- Precondition for defined behavior: i1 /= i2 branchMask :: Int -> Int -> Int -branchMask p1 p2 = - unsafeShiftL 1 (wordSize - 1 - countLeadingZeros (p1 `xor` p2)) +branchMask i1 i2 = iShiftRL (minBound :: Int) (countLeadingZeros (i1 `xor` i2)) {-# INLINE branchMask #-} + +-- | The shared prefix of two @Int@s. +-- +-- Precondition for defined behavior: i1 /= i2 +branchPrefix :: Int -> Int -> Prefix +branchPrefix i1 i2 = Prefix ((i1 .|. i2) .&. m) + where + m = unsafeShiftR (minBound :: Int) (countLeadingZeros (i1 `xor` i2)) +{-# INLINE branchPrefix #-} i2w :: Int -> Word i2w = fromIntegral diff --git a/src/Data/Map.hs b/src/Data/Map.hs --- a/src/Data/Map.hs +++ b/src/Data/Map.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Map @@ -20,7 +18,8 @@ -- This module re-exports the value lazy "Data.Map.Lazy" API. -- -- The @'Map' k v@ type represents a finite map (sometimes called a dictionary) --- from keys of type @k@ to values of type @v@. A 'Map' is strict in its keys but lazy +-- from keys of type @k@ to values of type @v@. Most operations require that @k@ +-- have an instance of the 'Ord' class. A 'Map' is strict in its keys but lazy -- in its values. -- -- The functions in "Data.Map.Strict" are careful to force values before @@ -32,7 +31,8 @@ -- * If you are using 'Prelude.Int' keys, you will get much better performance for most -- operations using "Data.IntMap.Lazy". -- --- * If you don't care about ordering, consider using @Data.HashMap.Lazy@ from the +-- * If you don't care about ordering and don't handle untrusted keys, consider +-- using @Data.HashMap.Lazy@ from the -- -- package instead. -- @@ -45,9 +45,11 @@ -- > import Data.Map (Map) -- > import qualified Data.Map as Map -- --- Note that the implementation is generally /left-biased/. Functions that take --- two maps as arguments and combine them, such as `union` and `intersection`, --- prefer the values in the first argument to those in the second. +-- The @'Ord' k@ instance is expected to be lawful and define a total order. +-- Unless otherwise specified, operations expect equality on keys to be +-- extensional: if keys @k1@ and @k2@ satisfy @k1 == k2@, they are considered +-- identical. For instance, if only one key must be retained by an operation, it +-- is free to select either. -- -- -- == Warning diff --git a/src/Data/Map/Internal.hs b/src/Data/Map/Internal.hs --- a/src/Data/Map/Internal.hs +++ b/src/Data/Map/Internal.hs @@ -1,19 +1,13 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} #if defined(__GLASGOW_HASKELL__) {-# LANGUAGE DeriveLift #-} {-# LANGUAGE RoleAnnotations #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE Trustworthy #-} {-# LANGUAGE TypeFamilies #-} -#define USE_MAGIC_PROXY 1 #endif -#ifdef USE_MAGIC_PROXY -{-# LANGUAGE MagicHash #-} -#endif - {-# OPTIONS_HADDOCK not-home #-} #include "containers.h" @@ -88,14 +82,11 @@ -- INLINABLE (that exposes the unfolding). --- [Note: Using INLINE] +-- 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. +-- We mark some functions INLINE where it is more beneficial than INLINABLE. +-- There is a price to pay -- code growth. The code INLINED is therefore usually +-- only the tree navigation, other work such as rebalancing is not INLINED. -- -- All methods marked INLINE have to be nonrecursive -- a 'go' function doing -- the real work is provided. @@ -159,10 +150,12 @@ -- ** Delete\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -202,6 +195,7 @@ , runWhenMissing , merge -- *** @WhenMatched@ tactics + , dropMatched , zipWithMaybeMatched , zipWithMatched -- *** @WhenMissing@ tactics @@ -231,6 +225,7 @@ , traverseMaybeMissing , traverseMissing , filterAMissing + , whenMissing -- ** Deprecated general combining function @@ -248,6 +243,7 @@ , mapKeys , mapKeysWith , mapKeysMonotonic + , mapAssocsMonotonic -- * Folds , foldr @@ -269,6 +265,9 @@ , keysSet , argSet , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA , fromArgSet -- ** Lists @@ -276,6 +275,7 @@ , fromList , fromListWith , fromListWithKey + , fromListUpsert -- ** Ordered lists , toAscList @@ -283,10 +283,12 @@ , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList , fromDescList , fromDescListWith , fromDescListWithKey + , fromDescListUpsert , fromDistinctDescList -- * Filter @@ -347,9 +349,6 @@ -- Used by the strict version , AreWeStrict (..) , atKeyImpl -#ifdef __GLASGOW_HASKELL__ - , atKeyPlain -#endif , bin , balance , balanceL @@ -363,7 +362,6 @@ , ascLinkAll , descLinkTop , descLinkAll - , MaybeS(..) , Identity(..) , Stack(..) , foldl'Stack @@ -401,8 +399,8 @@ import qualified Data.Set.Internal as Set import Data.Set.Internal (Set) import Utils.Containers.Internal.PtrEquality (ptrEq) -import Utils.Containers.Internal.StrictPair -import Utils.Containers.Internal.StrictMaybe +import Utils.Containers.Internal.Strict + (StrictPair(..), StrictTriple(..), toPair) import Utils.Containers.Internal.BitQueue import Utils.Containers.Internal.EqOrdUtil (EqM(..), OrdM(..)) #ifdef DEFINE_ALTERF_FALLBACK @@ -411,11 +409,12 @@ #if __GLASGOW_HASKELL__ import GHC.Exts (build, lazy) +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else import Language.Haskell.TH.Syntax (Lift) -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () -# ifdef USE_MAGIC_PROXY -import GHC.Exts (Proxy#, proxy# ) # endif import qualified GHC.Exts as GHCExts import Data.Data @@ -434,14 +433,14 @@ -- | \(O(\log n)\). Find the value at a key. -- Calls 'error' when the element can not be found. -- +-- __Note__: This function is partial. Prefer '!?'. +-- -- > 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__ {-# INLINE (!) #-} -#endif -- | \(O(\log n)\). Find the value at a key. -- Returns 'Nothing' when the element can not be found. @@ -453,16 +452,12 @@ (!?) :: Ord k => Map k a -> k -> Maybe a (!?) m k = lookup k m -#if __GLASGOW_HASKELL__ {-# INLINE (!?) #-} -#endif -- | Same as 'difference'. (\\) :: Ord k => Map k a -> Map k b -> Map k a m1 \\ m2 = difference m1 m2 -#if __GLASGOW_HASKELL__ {-# INLINE (\\) #-} -#endif {-------------------------------------------------------------------- Size balanced trees. @@ -488,7 +483,9 @@ instance (Ord k) => Monoid (Map k v) where mempty = empty mconcat = unions +#if !MIN_VERSION_base(4,11,0) mappend = (<>) +#endif -- | @(<>)@ = 'union' -- @@ -584,11 +581,7 @@ LT -> go k l GT -> go k r EQ -> Just x -#if __GLASGOW_HASKELL__ {-# INLINABLE lookup #-} -#else -{-# INLINE lookup #-} -#endif -- | \(O(\log n)\). Is the key a member of the map? See also 'notMember'. -- @@ -602,11 +595,7 @@ LT -> go k l GT -> go k r EQ -> True -#if __GLASGOW_HASKELL__ {-# INLINABLE member #-} -#else -{-# INLINE member #-} -#endif -- | \(O(\log n)\). Is the key not a member of the map? See also 'member'. -- @@ -615,14 +604,8 @@ notMember :: Ord k => k -> Map k a -> Bool notMember k m = not $ member k m -#if __GLASGOW_HASKELL__ {-# INLINABLE notMember #-} -#else -{-# INLINE notMember #-} -#endif --- | \(O(\log n)\). Find the value at a key. --- Calls 'error' when the element can not be found. find :: Ord k => k -> Map k a -> a find = go where @@ -631,11 +614,7 @@ LT -> go k l GT -> go k r EQ -> x -#if __GLASGOW_HASKELL__ {-# INLINABLE find #-} -#else -{-# INLINE find #-} -#endif -- | \(O(\log n)\). The expression @('findWithDefault' def k map)@ returns -- the value at key @k@ or returns default value @def@ @@ -651,11 +630,7 @@ LT -> go def k l GT -> go def k r EQ -> x -#if __GLASGOW_HASKELL__ {-# INLINABLE findWithDefault #-} -#else -{-# INLINE findWithDefault #-} -#endif -- | \(O(\log n)\). Find largest key smaller than the given one and return the -- corresponding (key, value) pair. @@ -672,11 +647,7 @@ goJust !_ kx' x' Tip = Just (kx', x') goJust k kx' x' (Bin _ kx x l r) | k <= kx = goJust k kx' x' l | otherwise = goJust k kx x r -#if __GLASGOW_HASKELL__ {-# INLINABLE lookupLT #-} -#else -{-# INLINE lookupLT #-} -#endif -- | \(O(\log n)\). Find smallest key greater than the given one and return the -- corresponding (key, value) pair. @@ -693,11 +664,7 @@ goJust !_ kx' x' Tip = Just (kx', x') goJust k kx' x' (Bin _ kx x l r) | k < kx = goJust k kx x l | otherwise = goJust k kx' x' r -#if __GLASGOW_HASKELL__ {-# INLINABLE lookupGT #-} -#else -{-# INLINE lookupGT #-} -#endif -- | \(O(\log n)\). Find largest key smaller or equal to the given one and return -- the corresponding (key, value) pair. @@ -717,11 +684,7 @@ goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx' x' l EQ -> Just (kx, x) GT -> goJust k kx x r -#if __GLASGOW_HASKELL__ {-# INLINABLE lookupLE #-} -#else -{-# INLINE lookupLE #-} -#endif -- | \(O(\log n)\). Find smallest key greater or equal to the given one and return -- the corresponding (key, value) pair. @@ -741,11 +704,7 @@ goJust k kx' x' (Bin _ kx x l r) = case compare k kx of LT -> goJust k kx x l EQ -> Just (kx, x) GT -> goJust k kx' x' r -#if __GLASGOW_HASKELL__ {-# INLINABLE lookupGE #-} -#else -{-# INLINE lookupGE #-} -#endif {-------------------------------------------------------------------- Construction @@ -801,11 +760,7 @@ where !r' = go orig kx x r EQ | x `ptrEq` y && (lazy orig `seq` (orig `ptrEq` ky)) -> t | otherwise -> Bin sz (lazy orig) x l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} -#else -{-# INLINE insert #-} -#endif #ifndef __GLASGOW_HASKELL__ lazy :: a -> a @@ -845,11 +800,7 @@ | otherwise -> balanceR ky y l r' where !r' = go orig kx x r EQ -> t -#if __GLASGOW_HASKELL__ {-# INLINABLE insertR #-} -#else -{-# INLINE insertR #-} -#endif -- | \(O(\log n)\). Insert with a function, combining new value and old value. -- @'insertWith' f key value mp@ @@ -878,11 +829,7 @@ GT -> balanceR ky y l (go f kx x r) EQ -> Bin sy kx (f x y) l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWith #-} -#else -{-# INLINE insertWith #-} -#endif -- | A helper function for 'unionWith'. When the key is already in -- the map, the key is left alone, not replaced. The combining @@ -901,11 +848,7 @@ LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> Bin sy ky (f y x) l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithR #-} -#else -{-# INLINE insertWithR #-} -#endif -- | \(O(\log n)\). Insert with a function, combining key, new value and old value. -- @'insertWithKey' f key value mp@ @@ -932,11 +875,7 @@ LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> Bin sy kx (f kx x y) l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKey #-} -#else -{-# INLINE insertWithKey #-} -#endif -- | A helper function for 'unionWithKey'. When the key is already in -- the map, the key is left alone, not replaced. The combining @@ -955,11 +894,7 @@ LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> Bin sy ky (f ky y x) l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKeyR #-} -#else -{-# INLINE insertWithKeyR #-} -#endif -- | \(O(\log n)\). Combines insert operation with old value retrieval. -- The expression (@'insertLookupWithKey' f k x map@) @@ -995,11 +930,7 @@ !t' = balanceR ky y l r' in (found :*: t') EQ -> (Just y :*: Bin sy kx (f kx x y) l r) -#if __GLASGOW_HASKELL__ {-# INLINABLE insertLookupWithKey #-} -#else -{-# INLINE insertLookupWithKey #-} -#endif {-------------------------------------------------------------------- Deletion @@ -1026,11 +957,55 @@ | otherwise -> balanceL kx x l r' where !r' = go k r EQ -> glue l r -#if __GLASGOW_HASKELL__ {-# INLINABLE delete #-} -#else -{-# INLINE delete #-} + +-- | \(O(\log n)\). Pop an entry from the map. +-- +-- Returns @Nothing@ if the key is not in the map. Otherwise returns @Just@ the +-- value at the key and a map with the entry removed. +-- +-- @ +-- pop 1 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Nothing +-- pop 2 (fromList [(0,"a"),(2,"b"),(4,"c")]) == Just ("b",fromList [(0,"a"),(4,"c")]) +-- @ +-- +-- @since 0.8.1 +pop :: Ord k => k -> Map k a -> Maybe (a, Map k a) +pop k0 t0 = case go k0 t0 of + Popped (Just y) t -> Just (y, t) + _ -> Nothing + where + -- See Note [Popped impl] + go !k (Bin _ kx x l r) = case compare k kx of + LT -> case go k l of + Popped y@(Just _) l' -> Popped y (balanceR kx x l' r) + q -> q + EQ -> Popped (Just x) (glue l r) + GT -> case go k r of + Popped y@(Just _) r' -> Popped y (balanceL kx x l r') + q -> q + go !_ Tip = Popped Nothing Tip +{-# INLINABLE pop #-} + +-- Note [Popped impl] +-- ~~~~~~~~~~~~~~~~~~ +-- Popped is implemented as a pair, though a sum makes more sense: +-- data Popped k a = NotPopped | Popped a !(Map k a) +-- This is because GHC optimizes a return value of `Popped k a` to +-- `(# Maybe a, Map k a #)`, avoiding all Popped allocations in `pop`. +-- GHC cannot do this with a sum type yet, see GHC #14259. Manually using +-- unboxed sums avoids the allocations but GHC loses strictness information, +-- see #25988. +-- +-- On GHC>=9.6 we unbox the Maybe and avoid that allocation too, so `pop`'s `go` +-- returns `(# (# (# #) | a #), Map k a #)`. + +data Popped k a = Popped +#if __GLASGOW_HASKELL__ >= 906 + {-# UNPACK #-} #endif + !(Maybe a) + !(Map k a) -- | \(O(\log n)\). Update a value at a specific key with the result of the provided function. -- When the key is not @@ -1042,11 +1017,7 @@ adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a adjust f = adjustWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ {-# INLINABLE adjust #-} -#else -{-# INLINE adjust #-} -#endif -- | \(O(\log n)\). Adjust a value at a specific key. When the key is not -- a member of the map, the original map is returned. @@ -1066,11 +1037,7 @@ LT -> Bin sx kx x (go f k l) r GT -> Bin sx kx x l (go f k r) EQ -> Bin sx kx (f kx x) l r -#if __GLASGOW_HASKELL__ {-# INLINABLE adjustWithKey #-} -#else -{-# INLINE adjustWithKey #-} -#endif -- | \(O(\log n)\). The expression (@'update' f k map@) updates the value @x@ -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is @@ -1083,11 +1050,7 @@ update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a update f = updateWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ {-# INLINABLE update #-} -#else -{-# INLINE update #-} -#endif -- | \(O(\log n)\). The expression (@'updateWithKey' f k map@) updates the -- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', @@ -1112,12 +1075,27 @@ EQ -> case f kx x of Just x' -> Bin sx kx x' l r Nothing -> glue l r -#if __GLASGOW_HASKELL__ {-# INLINABLE updateWithKey #-} -#else -{-# INLINE updateWithKey #-} -#endif +-- | \(O(\log n)\). Update the value at a key or insert a value if the key is +-- not in the map. +-- +-- @ +-- let inc = maybe 1 (+1) +-- upsert inc \'a\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',2),(\'c\',2)] +-- upsert inc \'b\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',1),(\'b\',1),(\'c\',2)] +-- @ +-- +-- @since 0.8.1 +upsert :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a +upsert f !k (Bin sz kx x l r) = + case compare k kx of + LT -> balanceL kx x (upsert f k l) r + EQ -> Bin sz kx (f (Just x)) l r + GT -> balanceR kx x l (upsert f k r) +upsert f !k Tip = singleton k (f Nothing) +{-# INLINABLE upsert #-} + -- | \(O(\log n)\). Look up and update. See also 'updateWithKey'. -- This function returns the changed value, if it is updated. -- Returns the original key value if the map entry is deleted. @@ -1126,6 +1104,8 @@ -- > 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 also: 'pop' -- See Note: Type of local 'go' function updateLookupWithKey :: Ord k => (k -> a -> Maybe a) -> k -> Map k a -> (Maybe a,Map k a) @@ -1145,11 +1125,7 @@ Just x' -> (Just x' :*: Bin sx kx x' l r) Nothing -> let !glued = glue l r in (Just x :*: glued) -#if __GLASGOW_HASKELL__ {-# INLINABLE updateLookupWithKey #-} -#else -{-# INLINE updateLookupWithKey #-} -#endif -- | \(O(\log n)\). The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. -- 'alter' can be used to insert, delete, or update a value in a 'Map'. @@ -1180,11 +1156,7 @@ EQ -> case f (Just x) of Just x' -> Bin sx kx x' l r Nothing -> glue l r -#if __GLASGOW_HASKELL__ {-# INLINABLE alter #-} -#else -{-# INLINE alter #-} -#endif -- Used to choose the appropriate alterF implementation. data AreWeStrict = Strict | Lazy @@ -1194,21 +1166,6 @@ -- or update a value in a 'Map'. In short: @'lookup' k \<$\> 'alterF' f k m = f -- ('lookup' k m)@. -- --- Example: --- --- @ --- interactiveAlter :: Int -> Map Int String -> IO (Map Int String) --- interactiveAlter k m = alterF f k m where --- f Nothing = do --- putStrLn $ show k ++ --- " was not found in the map. Would you like to add it?" --- getUserResponse1 :: IO (Maybe String) --- f (Just old) = do --- putStrLn $ "The key is currently bound to " ++ show old ++ --- ". Would you like to change or delete it?" --- getUserResponse2 :: IO (Maybe String) --- @ --- -- 'alterF' is the most general operation for working with an individual -- key that may or may not be in a given map. When used with trivial -- functors like 'Identity' and 'Const', it is often slightly slower than @@ -1229,15 +1186,27 @@ -- Note: 'alterF' is a flipped version of the @at@ combinator from -- @Control.Lens.At@. -- +-- === Examples +-- +-- @ +-- -- Lookup the value at the key, and also remove the existing value or set a new value. +-- lookupAndSet :: Ord k => k -> Maybe a -> Map k a -> (Maybe a, Map k a) +-- lookupAndSet k new = alterF (\\old -> (old, new)) k +-- @ +-- +-- @ +-- -- Delete the value at the key. If it is absent the result is Nothing. +-- mustDelete :: Ord k => k -> Map k a -> Maybe (Map k a) +-- mustDelete = alterF (Nothing <$) +-- @ +-- -- @since 0.5.8 alterF :: (Functor f, Ord k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a) alterF f k m = atKeyImpl Lazy k f m -#ifndef __GLASGOW_HASKELL__ -{-# INLINE alterF #-} -#else -{-# INLINABLE [2] alterF #-} +#ifdef __GLASGOW_HASKELL__ +{-# INLINE [2] alterF #-} -- We can save a little time by recognizing the special case of -- `Control.Applicative.Const` and just doing a lookup. @@ -1266,7 +1235,7 @@ case fres of Nothing -> case mv of Nothing -> m - Just old -> deleteAlong old q m + Just _ -> deleteAlong q m Just new -> case strict of Strict -> new `seq` case mv of Nothing -> insertAlong q k new m @@ -1290,7 +1259,14 @@ #endif #endif -data TraceResult a = TraceResult (Maybe a) {-# UNPACK #-} !BitQueue +-- On GHC >=9.6 we can unpack sum types, so we unbox the Maybe to avoid the +-- allocation. +data TraceResult a = TraceResult +#if __GLASGOW_HASKELL__ >= 906 + {-# UNPACK #-} +#endif + !(Maybe a) + {-# UNPACK #-} !BitQueue -- Look up a key and return a result indicating whether it was found -- and what path was taken. @@ -1303,12 +1279,7 @@ LT -> (go $! q `snocQB` False) k l GT -> (go $! q `snocQB` True) k r EQ -> TraceResult (Just x) (buildQ q) - -#ifdef __GLASGOW_HASKELL__ {-# INLINABLE lookupTrace #-} -#else -{-# INLINE lookupTrace #-} -#endif -- Insert at a location (which will always be a leaf) -- described by the path passed in. @@ -1322,47 +1293,13 @@ -- Delete from a location (which will always be a node) -- described by the path passed in. --- --- This is fairly horrifying! We don't actually have any --- use for the old value we're deleting. But if GHC sees --- that, then it will allocate a thunk representing the --- Map with the key deleted before we have any reason to --- believe we'll actually want that. This transformation --- enhances sharing, but we don't care enough about that. --- So deleteAlong needs to take the old value, and we need --- to convince GHC somehow that it actually uses it. We --- can't NOINLINE deleteAlong, because that would prevent --- the BitQueue from being unboxed. So instead we pass the --- old value to a NOINLINE constant function and then --- convince GHC that we use the result throughout the --- computation. Doing the obvious thing and just passing --- the value itself through the recursion costs 3-4% time, --- so instead we convert the value to a magical zero-width --- proxy that's ultimately erased. -deleteAlong :: any -> BitQueue -> Map k a -> Map k a -deleteAlong old !q0 !m = go (bogus old) q0 m where -#ifdef USE_MAGIC_PROXY - go :: Proxy# () -> BitQueue -> Map k a -> Map k a -#else - go :: any -> BitQueue -> Map k a -> Map k a -#endif - go !_ !_ Tip = Tip - go foom q (Bin _ ky y l r) = - case unconsQ q of - Just (False, tl) -> balanceR ky y (go foom tl l) r - Just (True, tl) -> balanceL ky y l (go foom tl r) - Nothing -> glue l r - -#ifdef USE_MAGIC_PROXY -{-# NOINLINE bogus #-} -bogus :: a -> Proxy# () -bogus _ = proxy# -#else --- No point hiding in this case. -{-# INLINE bogus #-} -bogus :: a -> a -bogus a = a -#endif +deleteAlong :: BitQueue -> Map k a -> Map k a +deleteAlong !_ Tip = Tip +deleteAlong !q (Bin _ ky y l r) = + case unconsQ q of + Just (False, tl) -> balanceR ky y (deleteAlong tl l) r + Just (True, tl) -> balanceL ky y l (deleteAlong tl r) + Nothing -> glue l r -- Replace the value found in the node described -- by the given path with a new one. @@ -1376,42 +1313,8 @@ #ifdef __GLASGOW_HASKELL__ atKeyIdentity :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> Map k a -> Identity (Map k a) -atKeyIdentity k f t = Identity $ atKeyPlain Lazy k (coerce f) t +atKeyIdentity k f t = Identity (alter (coerce f) k t) {-# INLINABLE atKeyIdentity #-} - -atKeyPlain :: Ord k => AreWeStrict -> k -> (Maybe a -> Maybe a) -> Map k a -> Map k a -atKeyPlain strict k0 f0 t = case go k0 f0 t of - AltSmaller t' -> t' - AltBigger t' -> t' - AltAdj t' -> t' - AltSame -> t - where - go :: Ord k => k -> (Maybe a -> Maybe a) -> Map k a -> Altered k a - go !k f Tip = case f Nothing of - Nothing -> AltSame - Just x -> case strict of - Lazy -> AltBigger $ singleton k x - Strict -> x `seq` (AltBigger $ singleton k x) - - go k f (Bin sx kx x l r) = case compare k kx of - LT -> case go k f l of - AltSmaller l' -> AltSmaller $ balanceR kx x l' r - AltBigger l' -> AltBigger $ balanceL kx x l' r - AltAdj l' -> AltAdj $ Bin sx kx x l' r - AltSame -> AltSame - GT -> case go k f r of - AltSmaller r' -> AltSmaller $ balanceL kx x l r' - AltBigger r' -> AltBigger $ balanceR kx x l r' - AltAdj r' -> AltAdj $ Bin sx kx x l r' - AltSame -> AltSame - EQ -> case f (Just x) of - Just x' -> case strict of - Lazy -> AltAdj $ Bin sx kx x' l r - Strict -> x' `seq` (AltAdj $ Bin sx kx x' l r) - Nothing -> AltSmaller $ glue l r -{-# INLINE atKeyPlain #-} - -data Altered k a = AltSmaller !(Map k a) | AltBigger !(Map k a) | AltAdj !(Map k a) | AltSame #endif #ifdef DEFINE_ALTERF_FALLBACK @@ -1454,6 +1357,8 @@ -- including, the 'size' of the map. Calls 'error' when the key is not -- a 'member' of the map. -- +-- __Note__: This function is partial. Prefer 'lookupIndex'. +-- -- > 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 @@ -1469,9 +1374,7 @@ LT -> go idx k l GT -> go (idx + size l + 1) k r EQ -> idx + size l -#if __GLASGOW_HASKELL__ {-# INLINABLE findIndex #-} -#endif -- | \(O(\log n)\). Look up 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 @@ -1492,14 +1395,14 @@ LT -> go idx k l GT -> go (idx + size l + 1) k r EQ -> Just $! idx + size l -#if __GLASGOW_HASKELL__ {-# INLINABLE lookupIndex #-} -#endif -- | \(O(\log n)\). Retrieve an element by its /index/, i.e. by its zero-based -- index in the sequence sorted by keys. If the /index/ is out of range (less -- than zero, greater or equal to 'size' of the map), 'error' is called. -- +-- __Note__: This function is partial. +-- -- > 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 @@ -1532,7 +1435,7 @@ go i (Bin _ kx x l r) = case compare i sizeL of LT -> go i l - GT -> link kx x l (go (i - sizeL - 1) r) + GT -> linkL kx x l (go (i - sizeL - 1) r) EQ -> l where sizeL = size l @@ -1552,7 +1455,7 @@ go !_ Tip = Tip go i (Bin _ kx x l r) = case compare i sizeL of - LT -> link kx x (go i l) r + LT -> linkR kx x (go i l) r GT -> go (i - sizeL - 1) r EQ -> insertMin kx x r where sizeL = size l @@ -1574,9 +1477,9 @@ go i (Bin _ kx x l r) = case compare i sizeL of LT -> case go i l of - ll :*: lr -> ll :*: link kx x lr r + ll :*: lr -> ll :*: linkR kx x lr r GT -> case go (i - sizeL - 1) r of - rl :*: rr -> link kx x l rl :*: rr + rl :*: rr -> linkL kx x l rl :*: rr EQ -> l :*: insertMin kx x r where sizeL = size l @@ -1584,6 +1487,8 @@ -- 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. -- +-- __Note__: This function is partial. +-- -- > 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 @@ -1610,6 +1515,8 @@ -- 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. -- +-- __Note__: This function is partial. +-- -- > 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 @@ -1667,6 +1574,8 @@ -- | \(O(\log n)\). The minimal key of the map. Calls 'error' if the map is empty. -- +-- __Note__: This function is partial. Prefer 'lookupMin'. +-- -- > findMin (fromList [(5,"a"), (3,"b")]) == (3,"b") -- > findMin empty Error: empty map has no minimal element @@ -1693,6 +1602,8 @@ -- | \(O(\log n)\). The maximal key of the map. Calls 'error' if the map is empty. -- +-- __Note__: This function is partial. Prefer 'lookupMax'. +-- -- > findMax (fromList [(5,"a"), (3,"b")]) == (5,"a") -- > findMax empty Error: empty map has no maximal element @@ -1832,9 +1743,7 @@ unions :: (Foldable f, Ord k) => f (Map k a) -> Map k a unions ts = Foldable.foldl' union empty ts -#if __GLASGOW_HASKELL__ -{-# INLINABLE unions #-} -#endif +{-# INLINE unions #-} -- Inline for list fusion -- | The union of a list of maps, with a combining operation: -- (@'unionsWith' f == 'Prelude.foldl' ('unionWith' f) 'empty'@). @@ -1845,9 +1754,7 @@ unionsWith :: (Foldable f, Ord k) => (a->a->a) -> f (Map k a) -> Map k a unionsWith f ts = Foldable.foldl' (unionWith f) empty ts -#if __GLASGOW_HASKELL__ {-# INLINABLE unionsWith #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). -- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and @t2@. @@ -1866,9 +1773,7 @@ | otherwise -> link k1 x1 l1l2 r1r2 where !l1l2 = union l1 l2 !r1r2 = union r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE union #-} -#endif {-------------------------------------------------------------------- Union with a combining function @@ -1891,9 +1796,7 @@ Just x2 -> link k1 (f x1 x2) l1l2 r1r2 where !l1l2 = unionWith f l1 l2 !r1r2 = unionWith f r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE unionWith #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). -- Union with a combining function. @@ -1914,9 +1817,7 @@ Just x2 -> link k1 (f k1 x1 x2) l1l2 r1r2 where !l1l2 = unionWithKey f l1 l2 !r1r2 = unionWithKey f r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE unionWithKey #-} -#endif {-------------------------------------------------------------------- Difference @@ -1943,9 +1844,7 @@ where !l1l2 = difference l1 l2 !r1r2 = difference r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE difference #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Remove all keys in a 'Set' from a 'Map'. -- @@ -1966,9 +1865,7 @@ where !lm' = withoutKeys lm ls !rm' = withoutKeys rm rs -#if __GLASGOW_HASKELL__ {-# INLINABLE withoutKeys #-} -#endif -- | \(O(n+m)\). Difference with a combining function. -- When two equal keys are @@ -1982,9 +1879,7 @@ differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a differenceWith f = merge preserveMissing dropMissing $ zipWithMaybeMatched (\_ x y -> f x y) -#if __GLASGOW_HASKELL__ {-# INLINABLE differenceWith #-} -#endif -- | \(O(n+m)\). Difference with a combining function. When two equal keys are -- encountered, the combining function is applied to the key and both values. @@ -1998,9 +1893,7 @@ differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a differenceWithKey f = merge preserveMissing dropMissing (zipWithMaybeMatched f) -#if __GLASGOW_HASKELL__ {-# INLINABLE differenceWithKey #-} -#endif {-------------------------------------------------------------------- @@ -2024,9 +1917,7 @@ !(l2, mb, r2) = splitMember k t2 !l1l2 = intersection l1 l2 !r1r2 = intersection r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE intersection #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Restrict a 'Map' to only those keys -- found in a 'Set'. @@ -2049,9 +1940,7 @@ !(l2, b, r2) = Set.splitMember k s !l1l2 = restrictKeys l1 l2 !r1r2 = restrictKeys r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE restrictKeys #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Intersection with a combining function. -- @@ -2069,9 +1958,7 @@ !(l2, mb, r2) = splitLookup k t2 !l1l2 = intersectionWith f l1 l2 !r1r2 = intersectionWith f r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE intersectionWith #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Intersection with a combining function. -- @@ -2088,9 +1975,7 @@ !(l2, mb, r2) = splitLookup k t2 !l1l2 = intersectionWithKey f l1 l2 !r1r2 = intersectionWithKey f r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE intersectionWithKey #-} -#endif {-------------------------------------------------------------------- Symmetric difference @@ -2120,9 +2005,7 @@ !(l2, found, r2) = splitMember k t2 !l1l2 = symmetricDifference l1 l2 !r1r2 = symmetricDifference r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE symmetricDifference #-} -#endif {-------------------------------------------------------------------- Disjoint @@ -2154,9 +2037,8 @@ {-------------------------------------------------------------------- Compose --------------------------------------------------------------------} --- | Relate the keys of one map to the values of --- the other, by using the values of the former as keys for lookups --- in the latter. +-- | Given maps @bc@ and @ab@, relate the keys of @ab@ to the values of @bc@, +-- by using the values of @ab@ as keys for lookups in @bc@. -- -- Complexity: \( O (n \log m) \), where \(m\) is the size of the first argument -- @@ -2205,13 +2087,12 @@ , missingKey :: k -> x -> f (Maybe y)} -- | @since 0.5.9 -instance (Applicative f, Monad f) => Functor (WhenMissing f k x) where +instance Monad f => Functor (WhenMissing f k x) where fmap = mapWhenMissing {-# INLINE fmap #-} -- | @since 0.5.9 -instance (Applicative f, Monad f) - => Category.Category (WhenMissing f k) where +instance Monad f => Category.Category (WhenMissing f k) where id = preserveMissing f . g = traverseMaybeMissing $ \ k x -> missingKey g k x >>= \y -> @@ -2224,7 +2105,7 @@ -- | Equivalent to @ ReaderT k (ReaderT x (MaybeT f)) @. -- -- @since 0.5.9 -instance (Applicative f, Monad f) => Applicative (WhenMissing f k x) where +instance Monad f => Applicative (WhenMissing f k x) where pure x = mapMissing (\ _ _ -> x) f <*> g = traverseMaybeMissing $ \k x -> do res1 <- missingKey f k x @@ -2237,7 +2118,7 @@ -- | Equivalent to @ ReaderT k (ReaderT x (MaybeT f)) @. -- -- @since 0.5.9 -instance (Applicative f, Monad f) => Monad (WhenMissing f k x) where +instance Monad f => Monad (WhenMissing f k x) where m >>= f = traverseMaybeMissing $ \k x -> do res1 <- missingKey m k x case res1 of @@ -2245,10 +2126,47 @@ Just r -> missingKey (f r) k x {-# INLINE (>>=) #-} +-- | Create a @WhenMissing@ from two functions. +-- +-- @whenMissing@ must be called with two functions @f@ and @g@ such that +-- @g = 'traverseMaybeWithKey' f@. @g@ may be a more efficient way of applying +-- @f@ to all key-value pairs in a @Map@. +-- +-- __Warning__: It is the caller's responsibility to ensure the above property. +-- +-- === __Examples__ +-- +-- @ +-- preserveMissing :: Applicative f => WhenMissing f k x x +-- preserveMissing = whenMissing f g +-- where +-- f _k x = pure (Just x) +-- g m = pure m +-- -- Note that this satisfies g = traverseMaybeWithKey f +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- For a usage of this, see examples on mergeA +-- isEmpty :: WhenMissing (Const All) k x y +-- isEmpty = whenMissing f g +-- where +-- f _k _x = Const (All False) +-- g m = Const (All (null m)) +-- -- Note that this satisfies g = traverseMaybeWithKey f +-- @ +-- +-- @since 0.8.1 +whenMissing + :: (k -> x -> f (Maybe y)) -> (Map k x -> f (Map k y)) -> WhenMissing f k x y +whenMissing = flip WhenMissing + -- | Map covariantly over a @'WhenMissing' f k x@. -- -- @since 0.5.9 -mapWhenMissing :: (Applicative f, Monad f) +mapWhenMissing :: Monad f => (a -> b) -> WhenMissing f k x a -> WhenMissing f k x b mapWhenMissing f t = WhenMissing @@ -2345,7 +2263,7 @@ {-# INLINE fmap #-} -- | @since 0.5.9 -instance (Monad f, Applicative f) => Category.Category (WhenMatched f k x) where +instance Monad f => Category.Category (WhenMatched f k x) where id = zipWithMatched (\_ _ y -> y) f . g = zipWithMaybeAMatched $ \k x y -> do @@ -2359,7 +2277,7 @@ -- | Equivalent to @ ReaderT k (ReaderT x (ReaderT y (MaybeT f))) @ -- -- @since 0.5.9 -instance (Monad f, Applicative f) => Applicative (WhenMatched f k x y) where +instance Monad f => Applicative (WhenMatched f k x y) where pure x = zipWithMatched (\_ _ _ -> x) fs <*> xs = zipWithMaybeAMatched $ \k x y -> do res <- runWhenMatched fs k x y @@ -2372,7 +2290,7 @@ -- | Equivalent to @ ReaderT k (ReaderT x (ReaderT y (MaybeT f))) @ -- -- @since 0.5.9 -instance (Monad f, Applicative f) => Monad (WhenMatched f k x y) where +instance Monad f => Monad (WhenMatched f k x y) where m >>= f = zipWithMaybeAMatched $ \k x y -> do res <- runWhenMatched m k x y case res of @@ -2398,6 +2316,13 @@ -- @since 0.5.9 type SimpleWhenMatched = WhenMatched Identity +-- | When a key is found in both maps, drop the key and values. +-- +-- @since 0.8.1 +dropMatched :: Applicative f => WhenMatched f k x y z +dropMatched = WhenMatched (\_ _ _ -> pure Nothing) +{-# INLINE dropMatched #-} + -- | When a key is found in both maps, apply a function to the -- key and values and use the result in the merged map. -- @@ -2609,7 +2534,7 @@ -- | Merge two maps. -- --- 'merge' takes two 'WhenMissing' tactics, a 'WhenMatched' +-- 'merge' takes two 'SimpleWhenMissing' tactics, a 'SimpleWhenMatched' -- tactic and two maps. It uses the tactics to merge the maps. -- Its behavior is best understood via its fundamental tactics, -- 'mapMaybeMissing' and 'zipWithMaybeMatched'. @@ -2617,45 +2542,31 @@ -- Consider -- -- @ --- merge (mapMaybeMissing g1) --- (mapMaybeMissing g2) --- (zipWithMaybeMatched f) --- m1 m2 +-- merge (mapMaybeMissing g1) (mapMaybeMissing g2) (zipWithMaybeMatched f) m1 m2 -- @ -- --- Take, for example, --- -- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- 'merge' will first \"align\" these maps by key: --- --- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- It will then pass the individual entries and pairs of entries --- to @g1@, @g2@, or @f@ as appropriate: --- --- @ --- maybes = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"] +-- g1 k x = if k == 2 then Just ("1" ++ x) else Nothing +-- g2 k x = if k == 3 then Just ("2" ++ x) else Nothing +-- f k x y = if k == 6 then Just ("3" ++ x ++ y) else Nothing +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- m2 = fromList [(3,"g"), (6,"h"), (9,"i"), (12,"j")] -- @ -- --- This produces a 'Maybe' for each key: +-- 'merge' will pass the keys and values to @g1@, @g2@, or @f@ as appropriate, +-- producing a @Maybe@ for each element. -- -- @ --- keys = 0 1 2 3 4 --- results = [Nothing, Just True, Just False, Nothing, Just True] +-- m1: [ (2, "a"), (4, "b"), (6, "c"), (8, "d"), (10, "e"), (12, "f")] +-- m2: [ (3, "g"), (6, "h"), (9, "i"), (12, "j")] +-- result: [ g1 2 "a", g2 3 "g", g1 4 "b", f 6 "c" "h", g1 8 "d", g2 9 "i", g1 10 "e", f 12 "f" "j"] +-- = [Just "1a", Just "2g", Nothing, Just "3ch", Nothing, Nothing, Nothing, Nothing] -- @ -- --- Finally, the @Just@ results are collected into a map: +-- The result map contains the @Just@ values. -- --- @ --- return value = [(1, True), (2, False), (4, True)] --- @ +-- >>> merge (mapMaybeMissing g1) (mapMaybeMissing g2) (zipWithMaybeMatched f) m1 m2 +-- fromList [(2,"1a"), (3,"2g"), (6,"3ch")] -- -- The other tactics below are optimizations or simplifications of -- 'mapMaybeMissing' for special cases. Most importantly, @@ -2684,7 +2595,7 @@ -> Map k a -- ^ Map @m1@ -> Map k b -- ^ Map @m2@ -> Map k c -merge g1 g2 f m1 m2 = runIdentity $ +merge g1 g2 f = \m1 m2 -> runIdentity $ mergeA g1 g2 f m1 m2 {-# INLINE merge #-} @@ -2695,49 +2606,44 @@ -- Its behavior is best understood via its fundamental tactics, -- 'traverseMaybeMissing' and 'zipWithMaybeAMatched'. -- +-- Behaves just like 'merge' while allowing @Applicative@ effects. Effects are +-- performed in increasing order of keys. +-- -- Consider -- -- @ -- mergeA (traverseMaybeMissing g1) --- (traverseMaybeMissing g2) --- (zipWithMaybeAMatched f) --- m1 m2 +-- (traverseMaybeMissing g2) +-- (zipWithMaybeAMatched f) +-- m1 +-- m2 -- @ -- --- Take, for example, --- -- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- @mergeA@ will first \"align\" these maps by key: --- --- @ --- m1 = [(0, \'a\'), (1, \'b\'), (3, \'c\'), (4, \'d\')] --- m2 = [(1, "one"), (2, "two"), (4, "three")] --- @ --- --- It will then pass the individual entries and pairs of entries --- to @g1@, @g2@, or @f@ as appropriate: --- --- @ --- actions = [g1 0 \'a\', f 1 \'b\' "one", g2 2 "two", g1 3 \'c\', f 4 \'d\' "three"] --- @ --- --- Next, it will perform the actions in the @actions@ list in order from --- left to right. --- --- @ --- keys = 0 1 2 3 4 --- results = [Nothing, Just True, Just False, Nothing, Just True] +-- g1 k x = let z = if k == 2 then Just ("1" ++ x) else Nothing +-- in z <$ putStrLn ("g1 " ++ show (k, x)) +-- g2 k x = let z = if k == 3 then Just ("2" ++ x) else Nothing +-- in z <$ putStrLn ("g2 " ++ show (k, x)) +-- f k x y = let z = if k == 6 then Just ("3" ++ x ++ y) else Nothing +-- in z <$ putStrLn ("f " ++ show (k, x, y)) +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- m2 = fromList [(3,"g"), (6,"h"), (9,"i"), (12,"j")] -- @ -- --- Finally, the @Just@ results are collected into a map: +-- As with 'merge', the result map is @[(2,"1a"), (3,"2g"), (6,"3ch")]@. +-- Additionally, @g1@, @g2@, and @f@ perform @IO@ effects, printing in +-- increasing order of key. -- --- @ --- return value = [(1, True), (2, False), (4, True)] --- @ +-- >>> mergeA (traverseMaybeMissing g1) (traverseMaybeMissing g2) (zipWithMaybeAMatched f) m1 m2 +-- g1 (2,"a") +-- g2 (3,"g") +-- g1 (4,"b") +-- f (6,"c","h") +-- g1 (8,"d") +-- g2 (9,"i") +-- g1 (10,"e") +-- f (12,"f","j") +-- fromList [(2,"1a"),(3,"2g"),(6,"3ch")] -- -- The other tactics below are optimizations or simplifications of -- 'traverseMaybeMissing' for special cases. Most importantly, @@ -2750,6 +2656,38 @@ -- site. To prevent excessive inlining, you should generally only use -- 'mergeA' to define custom combining functions. -- +-- === __Examples__ +-- +-- @ +-- data Pair a = Pair !a !a deriving Functor +-- +-- instance Applicative Pair where +-- pure x = Pair x x +-- liftA2 f (Pair x1 y1) (Pair x2 y2) = Pair (f x1 x2) (f y1 y2) +-- +-- -- | Calculate the left-biased union and intersection of two maps. +-- unionIntersection :: Ord k => Map k a -> Map k a -> (Map k a, Map k a) +-- unionIntersection m1 m2 = +-- case mergeA preserveAndDropMissing preserveAndDropMissing preserveLeftMatched m1 m2 of +-- Pair mu mi -> (mu, mi) +-- where +-- -- use Pair to build the union and intersection together +-- preserveAndDropMissing = 'whenMissing' (\\_k x -> Pair (Just x) Nothing) (\\m -> Pair m empty) +-- preserveLeftMatched = 'zipWithMaybeAMatched' (\\_k x1 _x2 -> Pair (Just x1) (Just x1)) +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- | Whether the keys of the first map are a subset of the keys of the second map. +-- keysAreSubsetOf :: Ord k => Map k a -> Map k b -> Bool +-- keysAreSubsetOf m1 m2 = +-- getAll (getConst (mergeA isEmpty 'dropMissing' 'dropMatched' m1 m2)) +-- where +-- isEmpty = 'whenMissing' (\\_k _x -> Const (All False)) (\\m -> Const (All (null m))) +-- @ +-- -- @since 0.5.9 mergeA :: (Applicative f, Ord k) @@ -2848,9 +2786,7 @@ -- isSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool isSubmapOf m1 m2 = isSubmapOfBy (==) m1 m2 -#if __GLASGOW_HASKELL__ {-# INLINABLE isSubmapOf #-} -#endif {- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if @@ -2875,9 +2811,7 @@ isSubmapOfBy :: Ord k => (a->b->Bool) -> Map k a -> Map k b -> Bool isSubmapOfBy f t1 t2 = size t1 <= size t2 && submap' f t1 t2 -#if __GLASGOW_HASKELL__ {-# INLINABLE isSubmapOfBy #-} -#endif -- Test whether a map is a submap of another without the *initial* -- size test. See Data.Set.Internal.isSubsetOfX for notes on @@ -2897,18 +2831,14 @@ && submap' f l lt && submap' f r gt where (lt,found,gt) = splitLookup kx t -#if __GLASGOW_HASKELL__ {-# INLINABLE submap' #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Is this a proper submap? (ie. a submap but not equal). -- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). isProperSubmapOf :: (Ord k,Eq a) => Map k a -> Map k a -> Bool isProperSubmapOf m1 m2 = isProperSubmapOfBy (==) m1 m2 -#if __GLASGOW_HASKELL__ {-# INLINABLE isProperSubmapOf #-} -#endif {- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Is this a proper submap? (ie. a submap but not equal). The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when @@ -2931,14 +2861,12 @@ isProperSubmapOfBy :: Ord k => (a -> b -> Bool) -> Map k a -> Map k b -> Bool isProperSubmapOfBy f t1 t2 = size t1 < size t2 && submap' f t1 t2 -#if __GLASGOW_HASKELL__ {-# INLINABLE isProperSubmapOfBy #-} -#endif {-------------------------------------------------------------------- Filter and partition --------------------------------------------------------------------} --- | \(O(n)\). Filter all values that satisfy the predicate. +-- | \(O(n)\). Keep all values that satisfy the predicate. -- -- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" -- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty @@ -2948,7 +2876,7 @@ filter p m = filterWithKey (\_ x -> p x) m --- | \(O(n)\). Filter all keys that satisfy the predicate. +-- | \(O(n)\). Keep all keys that satisfy the predicate. -- -- @ -- filterKeys p = 'filterWithKey' (\\k _ -> p k) @@ -2961,7 +2889,7 @@ filterKeys :: (k -> Bool) -> Map k a -> Map k a filterKeys p m = filterWithKey (\k _ -> p k) m --- | \(O(n)\). Filter all keys\/values that satisfy the predicate. +-- | \(O(n)\). Keep all keys\/values that satisfy the predicate. -- -- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" @@ -3001,7 +2929,7 @@ takeWhileAntitone :: (k -> Bool) -> Map k a -> Map k a takeWhileAntitone _ Tip = Tip takeWhileAntitone p (Bin _ kx x l r) - | p kx = link kx x l (takeWhileAntitone p r) + | p kx = linkL kx x l (takeWhileAntitone p r) | otherwise = takeWhileAntitone p l -- | \(O(\log n)\). Drop while a predicate on the keys holds. @@ -3019,7 +2947,7 @@ dropWhileAntitone _ Tip = Tip dropWhileAntitone p (Bin _ kx x l r) | p kx = dropWhileAntitone p r - | otherwise = link kx x (dropWhileAntitone p l) r + | otherwise = linkR kx x (dropWhileAntitone p l) r -- | \(O(\log n)\). Divide a map at the point where a predicate on the keys stops holding. -- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, @@ -3042,8 +2970,8 @@ where go _ Tip = Tip :*: Tip go p (Bin _ kx x l r) - | p kx = let u :*: v = go p r in link kx x l u :*: v - | otherwise = let u :*: v = go p l in u :*: link kx x v r + | p kx = let u :*: v = go p r in linkL kx x l u :*: v + | otherwise = let u :*: v = go p l in u :*: linkR kx x v r -- | \(O(n)\). Partition the map according to a predicate. The first -- map contains all elements that satisfy the predicate, the second all @@ -3114,6 +3042,7 @@ combine !l' mx !r' = case mx of Nothing -> link2 l' r' Just x' -> link kx x' l' r' +{-# INLINABLE traverseMaybeWithKey #-} -- | \(O(n)\). Map values and separate the 'Left' and 'Right' results. -- @@ -3262,9 +3191,7 @@ mapKeys :: Ord k2 => (k1->k2) -> Map k1 a -> Map k2 a mapKeys f m = finishB (foldlWithKey' (\b kx x -> insertB (f kx) x b) emptyB m) -#if __GLASGOW_HASKELL__ {-# INLINABLE mapKeys #-} -#endif -- | \(O(n \log n)\). -- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. @@ -3284,9 +3211,7 @@ mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a mapKeysWith c f m = finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB m) -#if __GLASGOW_HASKELL__ {-# INLINABLE mapKeysWith #-} -#endif -- | \(O(n)\). @@ -3311,10 +3236,26 @@ -- > 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) +mapKeysMonotonic f = mapAssocsMonotonic (\k x -> (f k, x)) +-- | \(O(n)\). Map over keys and values with a function @f@ that is +-- monotonically strictly increasing in the keys. That is, for keys @kx@ and +-- @ky@ and values @x@ and @y@, if @kx@ < @ky@ then +-- @fst (f kx x)@ < @fst (f ky y)@. +-- +-- __Warning__: This function should be used only if @f@ is monotonically +-- strictly increasing in the key. This precondition is not checked. +-- +-- @since 0.8.1 +mapAssocsMonotonic :: (k1 -> a1 -> (k2, a2)) -> Map k1 a1 -> Map k2 a2 +mapAssocsMonotonic f = go + where + go Tip = Tip + go (Bin sz k1 x1 l r) = case f k1 x1 of + (k2, x2) -> Bin sz k2 x2 (go l) (go r) +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE mapAssocsMonotonic #-} + {-------------------------------------------------------------------- Folds --------------------------------------------------------------------} @@ -3495,16 +3436,68 @@ argSet Tip = Set.Tip argSet (Bin sz kx x l r) = Set.Bin sz (Arg kx x) (argSet l) (argSet r) --- | \(O(n)\). Build a map from a set of keys and a function which for each key --- computes its value. +-- | \(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 -> replicate k 'a') (Data.Set.fromList [3,5]) == fromList [(3,"aaa"), (5,"aaaaa")] -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) +fromSet :: (k -> a) -> Set k -> Map k a +#ifdef __GLASGOW_HASKELL__ +fromSet = + (coerce :: ((k -> Identity a) -> Set k -> Identity (Map k a)) + -> (k -> a) -> Set k -> Map k a) + fromSetA +#else +fromSet f = runIdentity . fromSetA (pure . f) +#endif +-- | \(O(n)\). Build a map from a 'Set' of keys and a function which for each +-- key computes its value in an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- > let f k = if k == 0 then Nothing else Just (6 `div` k) +-- > fromSetA f (Data.Set.fromList [1,2,3,4]) == Just (fromList [(1,6),(2,3),(3,2),(4,1)]) +-- > fromSetA f (Data.Set.fromList [0,1,2]) == Nothing +-- +-- @since 0.8.1 +fromSetA :: Applicative f => (k -> f a) -> Set k -> f (Map k a) +fromSetA _ Set.Tip = pure Tip +fromSetA f (Set.Bin sz x l r) = + liftA3 (flip (Bin sz x)) (fromSetA f l) (f x) (fromSetA f r) +{-# INLINABLE fromSetA #-} + +-- | \(O(n)\). Build a map from a 'Set' of keys and a function which for each +-- key optionally computes its value. +-- +-- > let f k = if even k then Just (replicate k 'a') else Nothing +-- > fromSetMaybe f (Data.Set.fromList [1,2,3,4]) == fromList [(2,"aa"), (4,"aaaa")] +-- +-- @since 0.8.1 +fromSetMaybe :: (k -> Maybe a) -> Set k -> Map k a +#ifdef __GLASGOW_HASKELL__ +fromSetMaybe = + (coerce :: ((k -> Identity (Maybe a)) -> Set k -> Identity (Map k a)) + -> (k -> Maybe a) -> Set k -> Map k a) + fromSetMaybeA +#else +fromSetMaybe f s = runIdentity (fromSetMaybeA (Identity . f) s) +#endif + +-- | \(O(n)\). Build a map from a 'Set' of keys and a function which for each +-- key optionally computes its values in an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- @since 0.8.1 +fromSetMaybeA :: Applicative f => (k -> f (Maybe a)) -> Set k -> f (Map k a) +fromSetMaybeA f = go + where + go Set.Tip = pure Tip + go (Set.Bin _ k l r) = + liftA3 (\l' mx r' -> maybe link2 (link k) mx l' r') (go l) (f k) (go r) +{-# INLINABLE fromSetMaybeA #-} + -- | \(O(n)\). Build a map from a set of elements contained inside 'Arg's. -- -- > fromArgSet (Data.Set.fromList [Arg 3 "aaa", Arg 5 "aaaaa"]) == fromList [(5,"aaaaa"), (3,"aaa")] @@ -3527,7 +3520,7 @@ toList = toList #endif --- | \(O(n \log n)\). Build a map from a list of key\/value pairs. See also 'fromAscList'. +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs. -- If the list contains more than one value for the same key, the last value -- for the key is retained. -- @@ -3541,7 +3534,7 @@ fromList xs = finishB (Foldable.foldl' (\b (kx, x) -> insertB kx x b) emptyB xs) {-# INLINE fromList #-} -- INLINE for fusion --- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. -- -- If the keys are in non-decreasing order, this function takes \(O(n)\) time. -- @@ -3552,6 +3545,8 @@ -- -- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@. -- +-- See also: 'fromListUpsert' +-- -- === Performance -- -- You should ensure that the given @f@ is fast with this order of arguments. @@ -3584,7 +3579,7 @@ finishB (Foldable.foldl' (\b (kx, x) -> insertWithB f kx x b) emptyB xs) {-# INLINE fromListWith #-} -- INLINE for fusion --- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. -- -- If the keys are in non-decreasing order, this function takes \(O(n)\) time. -- @@ -3593,12 +3588,35 @@ -- > fromListWithKey f [] == empty -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromListUpsert' fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a fromListWithKey f xs = finishB (Foldable.foldl' (\b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs) {-# INLINE fromListWithKey #-} -- INLINE for fusion +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a +-- combining function. +-- +-- If the keys are in non-decreasing order, this function takes \(O(n)\) time. +-- +-- The result is equivalent to performing an @upsert@ for every key\/value in +-- the list. +-- +-- @ +-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty' +-- @ +-- +-- > let f x = maybe [x] (x:) +-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")] +-- +-- @since 0.8.1 +fromListUpsert :: Ord k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b +fromListUpsert f xs = + finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs) +{-# INLINE fromListUpsert #-} -- INLINE for fusion + -- | \(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")] @@ -3706,6 +3724,8 @@ -- > 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 +-- +-- See also: 'fromAscListUpsert' fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a fromAscListWith f xs @@ -3724,6 +3744,8 @@ -- -- Also see the performance note on 'fromListWith'. -- +-- See also: 'fromDescListUpsert' +-- -- @since 0.5.8 fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a @@ -3739,11 +3761,13 @@ -- if the precondition may not hold. -- -- > 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 f [(3,"b"), (5,"a"), (5,"b"), (5,"c")] == fromList [(3, "b"), (5, "5:c5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"c")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"c")]) == False -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next Nada xs) @@ -3769,6 +3793,8 @@ -- > valid (fromDescListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromDescListUpsert' fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a fromDescListWithKey f xs = descLinkAll (Foldable.foldl' next Nada xs) @@ -3781,7 +3807,50 @@ Nada -> Push ky y Tip stk {-# INLINE fromDescListWithKey #-} -- INLINE for fusion +-- | \(O(n)\). Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-decreasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromAscListUpsert f [(3,'a'), (3,'b'), (5,'c'), (5,'d'), (5,'e')] == fromList [(3,"ba"), (5,"edc")] +-- +-- @since 0.8.1 +fromAscListUpsert :: Eq k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b +fromAscListUpsert f xs = ascLinkAll (Foldable.foldl' next Nada xs) + where + next stk (!ky, y) = case stk of + Push kx x l stk' + | ky == kx -> Push ky (f y (Just x)) l stk' + | Tip <- l -> ascLinkTop stk' 1 (singleton kx x) ky (f y Nothing) + | otherwise -> Push ky (f y Nothing) Tip stk + Nada -> Push ky (f y Nothing) Tip stk +{-# INLINE fromAscListUpsert #-} -- INLINE for fusion +-- | \(O(n)\). Build a map from a descending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-increasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromDescListUpsert f [(5,'a'), (5,'b'), (5,'c'), (3,'d'), (3,'e')] == fromList [(3,"ed"), (5,"cba")] +-- +-- @since 0.8.1 +fromDescListUpsert :: Eq k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b +fromDescListUpsert f xs = descLinkAll (Foldable.foldl' next Nada xs) + where + next stk (!ky, y) = case stk of + Push kx x r stk' + | ky == kx -> Push ky (f y (Just x)) r stk' + | Tip <- r -> descLinkTop ky (f y Nothing) 1 (singleton kx x) stk' + | otherwise -> Push ky (f y Nothing) Tip stk + Nada -> Push ky (f y Nothing) Tip stk +{-# INLINE fromDescListUpsert #-} -- INLINE for fusion + -- | \(O(n)\). Build a map from an ascending list of distinct elements in linear time. -- -- __Warning__: This function should be used only if the keys are in @@ -3809,7 +3878,7 @@ ascLinkTop stk !_ l kx x = Push kx x l stk ascLinkAll :: Stack k a -> Map k a -ascLinkAll stk = foldl'Stack (\r kx x l -> link kx x l r) Tip stk +ascLinkAll stk = foldl'Stack (\r kx x l -> linkL kx x l r) Tip stk {-# INLINABLE ascLinkAll #-} -- | \(O(n)\). Build a map from a descending list of distinct elements in linear time. @@ -3842,7 +3911,7 @@ {-# INLINABLE descLinkTop #-} descLinkAll :: Stack k a -> Map k a -descLinkAll stk = foldl'Stack (\l kx x r -> link kx x l r) Tip stk +descLinkAll stk = foldl'Stack (\l kx x r -> linkR kx x l r) Tip stk {-# INLINABLE descLinkAll #-} data Stack k a = Push !k a !(Map k a) !(Stack k a) | Nada @@ -3906,12 +3975,10 @@ 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 + LT -> let (lt :*: gt) = go k l in lt :*: linkR kx x gt r + GT -> let (lt :*: gt) = go k r in linkL kx x l lt :*: gt EQ -> (l :*: r) -#if __GLASGOW_HASKELL__ {-# INLINABLE split #-} -#endif -- | \(O(\log n)\). The expression (@'splitLookup' k map@) splits a map just -- like 'split' but also returns @'lookup' k map@. @@ -3923,23 +3990,21 @@ -- > splitLookup 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], Nothing, empty) splitLookup :: Ord k => k -> Map k a -> (Map k a,Maybe a,Map k a) splitLookup k0 m = case go k0 m of - StrictTriple l mv r -> (l, mv, r) + TripleS l mv r -> (l, mv, r) where go :: Ord k => k -> Map k a -> StrictTriple (Map k a) (Maybe a) (Map k a) go !k t = case t of - Tip -> StrictTriple Tip Nothing Tip + Tip -> TripleS Tip Nothing Tip Bin _ kx x l r -> case compare k kx of - LT -> let StrictTriple lt z gt = go k l - !gt' = link kx x gt r - in StrictTriple lt z gt' - GT -> let StrictTriple lt z gt = go k r - !lt' = link kx x l lt - in StrictTriple lt' z gt - EQ -> StrictTriple l (Just x) r -#if __GLASGOW_HASKELL__ + LT -> let TripleS lt z gt = go k l + !gt' = linkR kx x gt r + in TripleS lt z gt' + GT -> let TripleS lt z gt = go k r + !lt' = linkL kx x l lt + in TripleS lt' z gt + EQ -> TripleS l (Just x) r {-# INLINABLE splitLookup #-} -#endif -- | \(O(\log n)\). A variant of 'splitLookup' that indicates only whether the -- key was present, rather than producing its value. This is used to @@ -3947,26 +4012,22 @@ -- constructors. splitMember :: Ord k => k -> Map k a -> (Map k a,Bool,Map k a) splitMember k0 m = case go k0 m of - StrictTriple l mv r -> (l, mv, r) + TripleS l mv r -> (l, mv, r) where go :: Ord k => k -> Map k a -> StrictTriple (Map k a) Bool (Map k a) go !k t = case t of - Tip -> StrictTriple Tip False Tip + Tip -> TripleS Tip False Tip Bin _ kx x l r -> case compare k kx of - LT -> let StrictTriple lt z gt = go k l - !gt' = link kx x gt r - in StrictTriple lt z gt' - GT -> let StrictTriple lt z gt = go k r - !lt' = link kx x l lt - in StrictTriple lt' z gt - EQ -> StrictTriple l True r -#if __GLASGOW_HASKELL__ + LT -> let TripleS lt z gt = go k l + !gt' = linkR kx x gt r + in TripleS lt z gt' + GT -> let TripleS lt z gt = go k r + !lt' = linkL kx x l lt + in TripleS lt' z gt + EQ -> TripleS l True r {-# INLINABLE splitMember #-} -#endif -data StrictTriple a b c = StrictTriple !a !b !c - {-------------------------------------------------------------------- MapBuilder --------------------------------------------------------------------} @@ -4012,6 +4073,21 @@ BMap m -> BMap (insertWith f ky y m) {-# INLINE insertWithB #-} +-- Upsert a key-value. The given function is used to generate the value based +-- on the existing value for the key. +upsertB :: Ord k => (Maybe a -> a) -> k -> MapBuilder k a -> MapBuilder k a +upsertB f !ky b = case b of + BAsc stk -> case stk of + Push kx x l stk' -> case compare ky kx of + LT -> BMap (upsert f ky (ascLinkAll stk)) + EQ -> BAsc (Push ky (f (Just x)) l stk') + GT -> case l of + Tip -> BAsc (ascLinkTop stk' 1 (singleton kx x) ky (f Nothing)) + Bin{} -> BAsc (Push ky (f Nothing) Tip stk) + Nada -> BAsc (Push ky (f Nothing) Tip Nada) + BMap m -> BMap (upsert f ky m) +{-# INLINE upsertB #-} + -- Finalize the builder into a Map. finishB :: MapBuilder k a -> Map k a finishB (BAsc stk) = ascLinkAll stk @@ -4046,12 +4122,39 @@ 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 +link kx x l@(Bin lsz lkx lx ll lr) r@(Bin rsz rkx rx rl rr) + | delta*lsz < rsz = balanceL rkx rx (linkR_ kx x lsz l rl) rr + | delta*rsz < lsz = balanceR lkx lx ll (linkL_ kx x lr rsz r) + | otherwise = Bin (1+lsz+rsz) kx x l r +-- Variant of link. Restores balance when the left tree may be too large for the +-- right tree, but not the other way around. +linkL :: k -> a -> Map k a -> Map k a -> Map k a +linkL kx x l r = case r of + Tip -> insertMax kx x l + Bin rsz _ _ _ _ -> linkL_ kx x l rsz r +linkL_ :: k -> a -> Map k a -> Int -> Map k a -> Map k a +linkL_ kx x l !rsz r = case l of + Bin lsz lkx lx ll lr + | delta*rsz < lsz -> balanceR lkx lx ll (linkL_ kx x lr rsz r) + | otherwise -> Bin (1+lsz+rsz) kx x l r + Tip -> Bin (1+rsz) kx x Tip r + +-- Variant of link. Restores balance when the right tree may be too large for +-- the left tree, but not the other way around. +linkR :: k -> a -> Map k a -> Map k a -> Map k a +linkR kx x l r = case l of + Tip -> insertMin kx x r + Bin lsz _ _ _ _ -> linkR_ kx x lsz l r + +linkR_ :: k -> a -> Int -> Map k a -> Map k a -> Map k a +linkR_ kx x !lsz l r = case r of + Bin rsz rkx rx rl rr + | delta*lsz < rsz -> balanceL rkx rx (linkR_ kx x lsz l rl) rr + | otherwise -> Bin (1+lsz+rsz) kx x l r + Tip -> Bin (1+lsz) kx x l Tip + -- insertMin and insertMax don't perform potentially expensive comparisons. insertMax,insertMin :: k -> a -> Map k a -> Map k a insertMax kx x t @@ -4072,11 +4175,25 @@ link2 :: Map k a -> Map k a -> Map k a link2 Tip r = r link2 l Tip = l -link2 l@(Bin sizeL kx x lx rx) r@(Bin sizeR ky y ly ry) - | delta*sizeL < sizeR = balanceL ky y (link2 l ly) ry - | delta*sizeR < sizeL = balanceR kx x lx (link2 rx r) - | otherwise = glue l r +link2 l@(Bin lsz lkx lx ll lr) r@(Bin rsz rkx rx rl rr) + | delta*lsz < rsz = balanceL rkx rx (link2R_ lsz l rl) rr + | delta*rsz < lsz = balanceR lkx lx ll (link2L_ lr rsz r) + | otherwise = glue l r +link2L_ :: Map k a -> Int -> Map k a -> Map k a +link2L_ l !rsz r = case l of + Bin lsz lkx lx ll lr + | delta*rsz < lsz -> balanceR lkx lx ll (link2L_ lr rsz r) + | otherwise -> glue l r + Tip -> r + +link2R_ :: Int -> Map k a -> Map k a -> Map k a +link2R_ !lsz l r = case r of + Bin rsz rkx rx rl rr + | delta*lsz < rsz -> balanceL rkx rx (link2R_ lsz l rl) rr + | otherwise -> glue l r + Tip -> l + {-------------------------------------------------------------------- [glue l r]: glues two trees together. Assumes that [l] and [r] are already balanced with respect to each other. @@ -4105,9 +4222,9 @@ -- | \(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 empty Error: can not return the minimal element of an empty map - +-- Calls 'error' if the map is empty. +-- +-- __Note__: This function is partial. Prefer 'minViewWithKey'. deleteFindMin :: Map k a -> ((k,a),Map k a) deleteFindMin t = case minViewWithKey t of Nothing -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip) @@ -4115,9 +4232,9 @@ -- | \(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 - +-- Calls 'error' if the map is empty. +-- +-- __Note__: This function is partial. Prefer 'maxViewWithKey'. deleteFindMax :: Map k a -> ((k,a),Map k a) deleteFindMax t = case maxViewWithKey t of Nothing -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip) @@ -4255,7 +4372,7 @@ (_, _) -> error "Failure in Data.Map.balance" {-# NOINLINE balance_ #-} --- Functions balanceL and balanceR are specialised versions of balance. +-- Functions balanceL and balanceR are specialized versions of balance. -- balanceL only checks whether the left subtree is too big, -- balanceR only checks whether the right subtree is too big. diff --git a/src/Data/Map/Internal/Debug.hs b/src/Data/Map/Internal/Debug.hs --- a/src/Data/Map/Internal/Debug.hs +++ b/src/Data/Map/Internal/Debug.hs @@ -1,6 +1,3 @@ -{-# LANGUAGE CPP #-} -#include "containers.h" - module Data.Map.Internal.Debug where import Data.Map.Internal (Map (..), size, delta) diff --git a/src/Data/Map/Lazy.hs b/src/Data/Map/Lazy.hs --- a/src/Data/Map/Lazy.hs +++ b/src/Data/Map/Lazy.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Map.Lazy @@ -18,7 +16,8 @@ -- = Finite Maps (lazy interface) -- -- The @'Map' k v@ type represents a finite map (sometimes called a dictionary) --- from keys of type @k@ to values of type @v@. A 'Map' is strict in its keys but lazy +-- from keys of type @k@ to values of type @v@. Most operations require that @k@ +-- have an instance of the 'Ord' class. A 'Map' is strict in its keys but lazy -- in its values. -- -- The functions in "Data.Map.Strict" are careful to force values before @@ -30,7 +29,8 @@ -- * If you are using 'Prelude.Int' keys, you will get much better performance for most -- operations using "Data.IntMap.Lazy". -- --- * If you don't care about ordering, consider using @Data.HashMap.Lazy@ from the +-- * If you don't care about ordering and don't handle untrusted keys, consider +-- using @Data.HashMap.Lazy@ from the -- -- package instead. -- @@ -43,9 +43,11 @@ -- > import Data.Map.Lazy (Map) -- > import qualified Data.Map.Lazy as Map -- --- Note that the implementation is generally /left-biased/. Functions that take --- two maps as arguments and combine them, such as `union` and `intersection`, --- prefer the values in the first argument to those in the second. +-- The @'Ord' k@ instance is expected to be lawful and define a total order. +-- Unless otherwise specified, operations expect equality on keys to be +-- extensional: if keys @k1@ and @k2@ satisfy @k1 == k2@, they are considered +-- identical. For instance, if only one key must be retained by an operation, it +-- is free to select either. -- -- -- == Warning @@ -105,26 +107,34 @@ -- * Construction , empty , singleton - , fromSet - , fromArgSet -- ** From Unordered Lists , fromList , fromListWith , fromListWithKey + , fromListUpsert -- ** From Ascending Lists , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList -- ** From Descending Lists , fromDescList , fromDescListWith , fromDescListWithKey + , fromDescListUpsert , fromDistinctDescList + -- ** From @Set@ + , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA + , fromArgSet + -- * Insertion , insert , insertWith @@ -133,10 +143,12 @@ -- * Deletion\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -206,6 +218,7 @@ , mapKeys , mapKeysWith , mapKeysMonotonic + , mapAssocsMonotonic -- * Folds , foldr @@ -272,12 +285,8 @@ -- * Min\/Max , lookupMin , lookupMax - , findMin - , findMax , deleteMin , deleteMax - , deleteFindMin - , deleteFindMax , updateMin , updateMax , updateMinWithKey @@ -286,6 +295,10 @@ , maxView , minViewWithKey , maxViewWithKey + , findMin + , findMax + , deleteFindMin + , deleteFindMax -- * Debugging , valid diff --git a/src/Data/Map/Merge/Lazy.hs b/src/Data/Map/Merge/Lazy.hs --- a/src/Data/Map/Merge/Lazy.hs +++ b/src/Data/Map/Merge/Lazy.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Map.Merge.Lazy @@ -42,6 +40,7 @@ , merge -- *** @WhenMatched@ tactics + , dropMatched , zipWithMaybeMatched , zipWithMatched @@ -73,6 +72,7 @@ , traverseMaybeMissing , traverseMissing , filterAMissing + , whenMissing -- *** Covariant maps for tactics , mapWhenMissing diff --git a/src/Data/Map/Merge/Set/Internal.hs b/src/Data/Map/Merge/Set/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Map/Merge/Set/Internal.hs @@ -0,0 +1,293 @@ +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- The contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- This module defines common constructs used by both "Data.Map.Merge.Set.Lazy" +-- and "Data.Map.Merge.Set.Strict". +-- +-- @since 0.8.1 +-- +module Data.Map.Merge.Set.Internal + ( WhenMatched(..) + , SimpleWhenMatched + , dropMatched + , filterMatched + , filterAMatched + + , WhenMissingSet(..) + , SimpleWhenMissingSet + , dropMissingSet + + , merge + , mergeA + + , runWhenMatched + , runWhenMissingSet + ) where + +import Control.Applicative (liftA3) +import Data.Functor.Identity (Identity(..)) +import Data.Set (Set) +import qualified Data.Set.Internal as S +import Data.Map (Map) +import qualified Data.Map.Internal as M + +-- | A tactic for dealing with keys present in both the set and the map in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @WhenMatched f k a b@ is an abstract representation of +-- a function of type @k -> a -> f (Maybe b)@. +-- +-- @since 0.8.1 +newtype WhenMatched f k a b = WhenMatched + { matchedKey :: k -> a -> f (Maybe b) + } + +-- | Run @WhenMatched@. +-- +-- @since 0.8.1 +runWhenMatched :: WhenMatched f k a b -> k -> a -> f (Maybe b) +runWhenMatched = matchedKey + +-- | A tactic for dealing with keys present in both the set and the map in +-- 'merge'. +-- +-- A tactic of type @SimpleWhenMatched k a b@ is an abstract representation of +-- a function of type @k -> a -> Maybe b@. +-- +-- @since 0.8.1 +type SimpleWhenMatched = WhenMatched Identity + +-- | When a key is found in both the map and the set, drop the key and value. +-- +-- @since 0.8.1 +dropMatched :: Applicative f => WhenMatched f k a b +dropMatched = WhenMatched (\_ _ -> pure Nothing) +{-# INLINE dropMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and keep the value in the merged map if the +-- result is @True@. +-- +-- @since 0.8.1 +filterMatched :: Applicative f => (k -> a -> Bool) -> WhenMatched f k a a +filterMatched f = + WhenMatched (\k x -> if f k x then pure (Just x) else pure Nothing) +{-# INLINE filterMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and keep the value in the merged map if the +-- result of the action is @True@. +-- +-- @since 0.8.1 +filterAMatched :: Functor f => (k -> a -> f Bool) -> WhenMatched f k a a +filterAMatched f = + WhenMatched (\k x -> (\b -> if b then Just x else Nothing) <$> f k x) +{-# INLINE filterAMatched #-} + +-- | A tactic for dealing with keys present in the set but not in the map in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @WhenMissingSet f k a@ is an abstract representation of +-- a function of type @k -> f (Maybe a)@. +-- +-- @since 0.8.1 +data WhenMissingSet f k a = WhenMissingSet + { missingSubtree :: Set k -> f (Map k a) + , missingKey :: k -> f (Maybe a) + } + +-- | Run @WhenMissingSet@. +-- +-- @since 0.8.1 +runWhenMissingSet :: WhenMissingSet f k a -> k -> f (Maybe a) +runWhenMissingSet = missingKey + +-- | A tactic for dealing with keys present in the set but not in the map in +-- 'merge'. +-- +-- A tactic of type @SimpleWhenMissingSet k a@ is an abstract representation of +-- a function of type @k -> Maybe a@. +-- +-- @since 0.8.1 +type SimpleWhenMissingSet = WhenMissingSet Identity + +-- | Drop keys that are present in the set but missing from the map. +-- +-- @since 0.8.1 +dropMissingSet :: Applicative f => WhenMissingSet f k a +dropMissingSet = WhenMissingSet + { missingSubtree = \_ -> pure M.empty + , missingKey = \_ -> pure Nothing + } +{-# INLINE dropMissingSet #-} + +-- | Merge a map and a set into a map. +-- +-- 'merge' takes a 'M.SimpleWhenMissing' tactic, a 'SimpleWhenMissingSet' +-- tactic, a 'SimpleWhenMatched' tactic, a map and a set. It uses the tactics to +-- merge the map and the set into a map. +-- +-- Its behavior is best understood via the tactics @mapMaybeMissing@, +-- @generateMaybeMissingSet@, and @mapMaybeMatched@. Consider +-- +-- @ +-- merge (mapMaybeMissing g1) (generateMaybeMissingSet g2) (mapMaybeMatched f) m1 s2 +-- @ +-- +-- @ +-- g1 k x = if k == 2 then Just ("1" ++ x) else Nothing +-- g2 k = if k == 3 then Just "2" else Nothing +-- f k x = if k == 6 then Just ("3" ++ x) else Nothing +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- s2 = fromList [3, 6, 9, 12] +-- @ +-- +-- 'merge' will pass the keys and values to @g1@, @g2@, or @f@ as appropriate, +-- producing a @Maybe@ for each element. +-- +-- @ +-- m1: [ (2, "a"), (4, "b"), (6, "c"), (8, "d"), (10, "e"), (12, "f")] +-- s2: [ 3, 6, 9, 12] +-- result: [ g1 2 "a", g2 3, g1 4 "b", f 6 "c", g1 8 "d", g2 9, g1 10 "e", f 12 "f"] +-- = [Just "1a", Just "2", Nothing, Just "3c", Nothing, Nothing, Nothing, Nothing] +-- @ +-- +-- The result map contains the @Just@ values. +-- +-- >>> merge (mapMaybeMissing g1) (generateMaybeMissingSet g2) (mapMaybeMatched f) m1 s2 +-- fromList [(2,"1a"), (3,"2g"), (6,"3c")] +-- +-- When 'merge' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should typically use 'merge' +-- to define your custom combining functions. +-- +-- @since 0.8.1 +merge + :: Ord k + => M.SimpleWhenMissing k a b -- ^ What to do with keys in @m1@ but not @s2@ + -> SimpleWhenMissingSet k b -- ^ What to do with keys in @s2@ but not @m1@ + -> SimpleWhenMatched k a b -- ^ What to do with keys in both @m1@ and @s2@ + -> Map k a -- ^ Map @m1@ + -> Set k -- ^ Set @s2@ + -> Map k b +merge miss1 miss2 match = \t1 t2 -> runIdentity (mergeA miss1 miss2 match t1 t2) +{-# INLINE merge #-} + +-- | Merge a map and a set into a map. Applicative version of 'merge'. +-- +-- 'mergeA' takes a 'M.WhenMissing' tactic, a 'WhenMissingSet' tactic, a +-- 'WhenMatched' tactic, a map and a set. It uses the tactics to merge the map +-- and the set into a map. +-- +-- Behaves just like 'merge' while allowing @Applicative@ effects. Effects are +-- performed in increasing order of keys. +-- +-- Consider +-- +-- @ +-- mergeA (traverseMaybeMissing g1) +-- (generateMaybeAMissingSet g2) +-- (traverseMaybeMatched f) +-- m1 +-- s2 +-- @ +-- +-- @ +-- g1 k x = let z = if k == 2 then Just ("1" ++ x) else Nothing +-- in z <$ putStrLn ("g1 " ++ show (k, x)) +-- g2 k = let z = if k == 3 then Just "2" else Nothing +-- in z <$ putStrLn ("g2 " ++ show k) +-- f k x = let z = if k == 6 then Just ("3" ++ x) else Nothing +-- in z <$ putStrLn ("f " ++ show (k, x)) +-- m1 = fromList [(2,"a"), (4,"b"), (6,"c"), (8,"d"), (10,"e"), (12,"f")] +-- m2 = fromList [3, 6, 9, 12] +-- @ +-- +-- As with 'merge', the result map is @[(2,"1a"), (3,"2"), (6,"3c")]@. +-- Additionally, @g1@, @g2@, and @f@ perform @IO@ effects, printing in +-- increasing order of key. +-- +-- >>> mergeA (traverseMaybeMissing g1) (generateMaybeAMissingSet g2) (traverseMaybeMatched f) m1 s2 +-- g1 (2,"a") +-- g2 3 +-- g1 (4,"b") +-- f (6,"c") +-- g1 (8,"d") +-- g2 9 +-- g1 (10,"e") +-- f (12,"f") +-- fromList [(2,"1a"),(3,"2"),(6,"3c")] +-- +-- When 'mergeA' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should generally only use +-- 'mergeA' to define custom combining functions. +-- +-- === __Examples__ +-- +-- @ +-- data Pair a = Pair !a !a deriving Functor +-- +-- instance Applicative Pair where +-- pure x = Pair x x +-- liftA2 f (Pair x1 y1) (Pair x2 y2) = Pair (f x1 x2) (f y1 y2) +-- +-- -- | Partition the map according to whether the keys appear in the set. +-- partitionKeys :: Ord k => Map k a -> Set k -> (Map k a, Map k a) +-- partitionKeys m s = +-- case mergeA dropAndPreserveMissing dropMissingSet preserveAndDropMatched m s of +-- Pair m1 m2 -> (m1, m2) +-- where +-- dropAndPreserveMissing = whenMissing (\\_k x -> Pair Nothing (Just x)) (\\m -> Pair empty m) +-- preserveAndDropMatched = traverseMaybeMatched (\\_k x -> Pair (Just x) Nothing) +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- | Whether the keys of the map are a subset of the keys of the set. +-- keysAreSubsetOf :: Ord k => Map k a -> Set k -> Bool +-- keysAreSubsetOf m s = +-- getAll (getConst (mergeA isEmpty dropMissing 'dropMatched' m1 m2)) +-- where +-- isEmpty = whenMissing (\\_k _x -> Const (All False)) (\\m -> Const (All (null m))) +-- @ +-- +-- @since 0.8.1 +mergeA + :: (Applicative f, Ord k) + => M.WhenMissing f k a b -- ^ What to do with keys in @m1@ but not @s2@ + -> WhenMissingSet f k b -- ^ What to do with keys in @s2@ but not @m1@ + -> WhenMatched f k a b -- ^ What to do with keys in both @m1@ and @s2@ + -> Map k a -- ^ Map @m1@ + -> Set k -- ^ Set @s2@ + -> f (Map k b) +mergeA + M.WhenMissing{M.missingSubtree = g1t, M.missingKey = g1k} + WhenMissingSet{missingSubtree = g2t} + WhenMatched{matchedKey = f} = go + where + go t1 S.Tip = g1t t1 + go M.Tip t2 = g2t t2 + go (M.Bin _ k1 x1 l1 r1) t2 = case S.splitMember k1 t2 of + (l2, found, r2) -> + liftA3 + (\l' mx' r' -> maybe M.link2 (M.link k1) mx' l' r') + (go l1 l2) + (if found then f k1 x1 else g1k k1 x1) + (go r1 r2) +{-# INLINE mergeA #-} diff --git a/src/Data/Map/Merge/Set/Lazy.hs b/src/Data/Map/Merge/Set/Lazy.hs new file mode 100644 --- /dev/null +++ b/src/Data/Map/Merge/Set/Lazy.hs @@ -0,0 +1,155 @@ +-- | +-- This module defines an API for writing functions that merge a map and a +-- set into a map. The key functions are 'Data.Map.Merge.Set.Lazy.merge' and +-- 'Data.Map.Merge.Set.Lazy.mergeA'. Each of these can be used with several +-- different \"merge tactics\". +-- +-- The @merge@ and @mergeA@ functions are shared by the lazy and strict +-- modules. Only the choice of merge tactics determines strictness. If you +-- use 'Data.Map.Merge.Set.Strict.mapMissing' from "Data.Map.Merge.Set.Strict" +-- then the results will be forced before they are inserted. If you use +-- 'Data.Map.Merge.Set.Lazy.mapMissing' from this module then they will not. +-- +-- @since 0.8.1 +-- +module Data.Map.Merge.Set.Lazy + ( + -- ** Simple merge tactic types + M.SimpleWhenMissing + , Internal.SimpleWhenMissingSet + , Internal.SimpleWhenMatched + + -- ** General combining function + , Internal.merge + + -- *** @WhenMatched@ tactics + , Internal.dropMatched + , Internal.filterMatched + , mapMatched + , mapMaybeMatched + + -- *** @WhenMissing@ tactics + , M.dropMissing + , M.preserveMissing + , M.mapMissing + , M.filterMissing + , M.mapMaybeMissing + + -- *** @WhenMissingSet@ tactics + , Internal.dropMissingSet + , generateMissingSet + , generateMaybeMissingSet + + -- ** Applicative merge tactic types + , M.WhenMissing + , Internal.WhenMissingSet + , Internal.WhenMatched + + -- ** General combining function + , Internal.mergeA + + -- *** @WhenMatched@ tactics + , Internal.filterAMatched + , traverseMatched + , traverseMaybeMatched + + -- *** @WhenMissing@ tactics + , M.filterAMissing + , M.traverseMissing + , M.traverseMaybeMissing + , M.whenMissing + + -- *** @WhenMissingSet@ tactics + , generateAMissingSet + , generateMaybeAMissingSet + + -- ** Miscellaneous + , Internal.runWhenMatched + , Internal.runWhenMissingSet + ) where + +import qualified Data.Map.Internal as M +import qualified Data.Map.Merge.Set.Internal as Internal +import Data.Map.Merge.Set.Internal (WhenMatched(..), WhenMissingSet(..)) + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and use the result as the value for the merged +-- map. +-- +-- @since 0.8.1 +mapMatched :: Applicative f => (k -> a -> b) -> WhenMatched f k a b +mapMatched f = WhenMatched (\k x -> pure (Just (f k x))) +{-# INLINE mapMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and maybe use the result as the value for the +-- merged map. +-- +-- @since 0.8.1 +mapMaybeMatched :: Applicative f => (k -> a -> Maybe b) -> WhenMatched f k a b +mapMaybeMatched f = WhenMatched (\k x -> pure (f k x)) +{-# INLINE mapMaybeMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and use the result of the action as the value +-- for the merged map. +-- +-- @since 0.8.1 +traverseMatched :: Functor f => (k -> a -> f b) -> WhenMatched f k a b +traverseMatched f = WhenMatched (\k x -> Just <$> f k x) +{-# INLINE traverseMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and maybe use the result of the action as the +-- value for the merged map. +-- +-- @since 0.8.1 +traverseMaybeMatched :: (k -> a -> f (Maybe b)) -> WhenMatched f k a b +traverseMaybeMatched = WhenMatched + +-- | For keys that are present in the set but missing from the map, apply a +-- function and use the result as the value for the merged map. +-- +-- @since 0.8.1 +generateMissingSet :: Applicative f => (k -> a) -> WhenMissingSet f k a +generateMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (M.fromSet f s) + , missingKey = \k -> pure (Just (f k)) + } +{-# INLINE generateMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and use the result of the action as the value for the merged map. +-- +-- @since 0.8.1 +generateAMissingSet :: Applicative f => (k -> f a) -> WhenMissingSet f k a +generateAMissingSet f = WhenMissingSet + { missingSubtree = M.fromSetA f + , missingKey = \k -> Just <$> f k + } +{-# INLINE generateAMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function and maybe use the result as the value for the merged map. +-- +-- @since 0.8.1 +generateMaybeMissingSet + :: Applicative f => (k -> Maybe a) -> WhenMissingSet f k a +generateMaybeMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (M.fromSetMaybe f s) + , missingKey = \k -> pure (f k) + } +{-# INLINE generateMaybeMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and maybe use the result of the action as the value for the merged +-- map. +-- +-- @since 0.8.1 +generateMaybeAMissingSet + :: Applicative f => (k -> f (Maybe a)) -> WhenMissingSet f k a +generateMaybeAMissingSet f = WhenMissingSet + { missingSubtree = M.fromSetMaybeA f + , missingKey = f + } +{-# INLINE generateMaybeAMissingSet #-} diff --git a/src/Data/Map/Merge/Set/Strict.hs b/src/Data/Map/Merge/Set/Strict.hs new file mode 100644 --- /dev/null +++ b/src/Data/Map/Merge/Set/Strict.hs @@ -0,0 +1,165 @@ +{-# LANGUAGE BangPatterns #-} + +-- | +-- This module defines an API for writing functions that merge a map and a +-- set into a map. The key functions are 'Data.Map.Merge.Set.Strict.merge' and +-- 'Data.Map.Merge.Set.Strict.mergeA'. Each of these can be used with several +-- different \"merge tactics\". +-- +-- The @merge@ and @mergeA@ functions are shared by the lazy and strict +-- modules. Only the choice of merge tactics determines strictness. +-- If you use 'Data.Map.Merge.Set.Strict.mapMissing' from this module +-- then the results will be forced before they are inserted. If you use +-- 'Data.Map.Merge.Set.Lazy.mapMissing' from "Data.Map.Merge.Set.Lazy" then they +-- will not. +-- +-- @since 0.8.1 +-- +module Data.Map.Merge.Set.Strict + ( + -- ** Simple merge tactic types + MS.SimpleWhenMissing + , Internal.SimpleWhenMissingSet + , Internal.SimpleWhenMatched + + -- ** General combining function + , Internal.merge + + -- *** @WhenMatched@ tactics + , Internal.dropMatched + , Internal.filterMatched + , mapMatched + , mapMaybeMatched + + -- *** @WhenMissing@ tactics + , MS.dropMissing + , MS.preserveMissing + , MS.mapMissing + , MS.filterMissing + , MS.mapMaybeMissing + + -- *** @WhenMissingSet@ tactics + , Internal.dropMissingSet + , generateMissingSet + , generateMaybeMissingSet + + -- ** Applicative merge tactic types + , MS.WhenMissing + , Internal.WhenMissingSet + , Internal.WhenMatched + + -- ** General combining function + , Internal.mergeA + + -- *** @WhenMatched@ tactics + , Internal.filterAMatched + , traverseMatched + , traverseMaybeMatched + + -- *** @WhenMissing@ tactics + , MS.filterAMissing + , MS.traverseMissing + , MS.traverseMaybeMissing + , M.whenMissing + + -- *** @WhenMissingSet@ tactics + , generateAMissingSet + , generateMaybeAMissingSet + + -- ** Miscellaneous + , Internal.runWhenMatched + , Internal.runWhenMissingSet + ) where + +import qualified Data.Map.Strict.Internal as MS +import qualified Data.Map.Internal as M +import qualified Data.Map.Merge.Set.Internal as Internal +import Data.Map.Merge.Set.Internal (WhenMatched(..), WhenMissingSet(..)) + +-- | When the key is found in both the map and the set, apply a function to the +-- key and the value in the map and use the result as the value for the result +-- map. +-- +-- @since 0.8.1 +mapMatched :: Applicative f => (k -> a -> b) -> WhenMatched f k a b +mapMatched f = WhenMatched (\k x -> pure (Just $! f k x)) +{-# INLINE mapMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map and maybe use the result as the value for the +-- merged map. +-- +-- @since 0.8.1 +mapMaybeMatched :: Applicative f => (k -> a -> Maybe b) -> WhenMatched f k a b +mapMaybeMatched f = WhenMatched (\k x -> pure (forceMaybe (f k x))) +{-# INLINE mapMaybeMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and use the result of the action as the value +-- for the merged map. +-- +-- @since 0.8.1 +traverseMatched :: Functor f => (k -> a -> f b) -> WhenMatched f k a b +traverseMatched f = WhenMatched (\k x -> (Just $!) <$> f k x) +{-# INLINE traverseMatched #-} + +-- | When a key is found in both the map and the set, apply a function to the +-- key and the value in the map, and maybe use the result of the action as the +-- value for the merged map. +-- +-- @since 0.8.1 +traverseMaybeMatched + :: Functor f => (k -> a -> f (Maybe b)) -> WhenMatched f k a b +traverseMaybeMatched f = WhenMatched (\k x -> forceMaybe <$> f k x) +{-# INLINE traverseMaybeMatched #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function and use the result as the value for the merge map. +-- +-- @since 0.8.1 +generateMissingSet :: Applicative f => (k -> a) -> WhenMissingSet f k a +generateMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (MS.fromSet f s) + , missingKey = \k -> pure (Just $! f k) + } +{-# INLINE generateMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and use the result of the action as the value for the merged map. +-- +-- @since 0.8.1 +generateAMissingSet :: Applicative f => (k -> f a) -> WhenMissingSet f k a +generateAMissingSet f = WhenMissingSet + { missingSubtree = MS.fromSetA f + , missingKey = \k -> (Just $!) <$> f k + } +{-# INLINE generateAMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function and maybe use the result as the value for the merged map. +-- +-- @since 0.8.1 +generateMaybeMissingSet + :: Applicative f => (k -> Maybe a) -> WhenMissingSet f k a +generateMaybeMissingSet f = WhenMissingSet + { missingSubtree = \s -> pure (MS.fromSetMaybe f s) + , missingKey = \k -> pure (f k) + } +{-# INLINE generateMaybeMissingSet #-} + +-- | For keys that are present in the set but missing from the map, apply a +-- function, and maybe use the result of the action as the value for the merged +-- map. +-- +-- @since 0.8.1 +generateMaybeAMissingSet + :: Applicative f => (k -> f (Maybe a)) -> WhenMissingSet f k a +generateMaybeAMissingSet f = WhenMissingSet + { missingSubtree = MS.fromSetMaybeA f + , missingKey = \k -> forceMaybe <$> f k + } +{-# INLINE generateMaybeAMissingSet #-} + +forceMaybe :: Maybe a -> Maybe a +forceMaybe Nothing = Nothing +forceMaybe m@(Just !_) = m diff --git a/src/Data/Map/Merge/Strict.hs b/src/Data/Map/Merge/Strict.hs --- a/src/Data/Map/Merge/Strict.hs +++ b/src/Data/Map/Merge/Strict.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Map.Merge.Strict @@ -47,6 +45,7 @@ , merge -- *** @WhenMatched@ tactics + , dropMatched , zipWithMaybeMatched , zipWithMatched @@ -79,6 +78,7 @@ , traverseMaybeMissing , traverseMissing , filterAMissing + , Internal.whenMissing -- ** Covariant maps for tactics , mapWhenMissing @@ -90,4 +90,5 @@ , runWhenMissing ) where +import qualified Data.Map.Internal as Internal import Data.Map.Strict.Internal diff --git a/src/Data/Map/Strict.hs b/src/Data/Map/Strict.hs --- a/src/Data/Map/Strict.hs +++ b/src/Data/Map/Strict.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Map.Strict @@ -18,7 +16,8 @@ -- = Finite Maps (strict interface) -- -- The @'Map' k v@ type represents a finite map (sometimes called a dictionary) --- from keys of type @k@ to values of type @v@. +-- from keys of type @k@ to values of type @v@. Most operations require that @k@ +-- have an instance of the 'Ord' class. -- -- Each function in this module is careful to force values before installing -- them in a 'Map'. This is usually more efficient when laziness is not @@ -35,7 +34,8 @@ -- * If you are using 'Prelude.Int' keys, you will get much better performance for -- most operations using "Data.IntMap.Strict". -- --- * If you don't care about ordering, consider use @Data.HashMap.Strict@ from the +-- * If you don't care about ordering and don't handle untrusted keys, consider +-- using @Data.HashMap.Strict@ from the -- -- package instead. -- @@ -48,9 +48,11 @@ -- > import Data.Map.Strict (Map) -- > import qualified Data.Map.Strict as Map -- --- Note that the implementation is generally /left-biased/. Functions that take --- two maps as arguments and combine them, such as `union` and `intersection`, --- prefer the values in the first argument to those in the second. +-- The @'Ord' k@ instance is expected to be lawful and define a total order. +-- Unless otherwise specified, operations expect equality on keys to be +-- extensional: if keys @k1@ and @k2@ satisfy @k1 == k2@, they are considered +-- identical. For instance, if only one key must be retained by an operation, it +-- is free to select either. -- -- -- == Warning @@ -119,26 +121,34 @@ -- * Construction , empty , singleton - , fromSet - , fromArgSet -- ** From Unordered Lists , fromList , fromListWith , fromListWithKey + , fromListUpsert -- ** From Ascending Lists , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList -- ** From Descending Lists , fromDescList , fromDescListWith , fromDescListWithKey + , fromDescListUpsert , fromDistinctDescList + -- ** From @Set@ + , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA + , fromArgSet + -- * Insertion , insert , insertWith @@ -147,10 +157,12 @@ -- * Deletion\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -220,6 +232,7 @@ , mapKeys , mapKeysWith , mapKeysMonotonic + , mapAssocsMonotonic -- * Folds , foldr @@ -287,12 +300,8 @@ -- * Min\/Max , lookupMin , lookupMax - , findMin - , findMax , deleteMin , deleteMax - , deleteFindMin - , deleteFindMax , updateMin , updateMax , updateMinWithKey @@ -301,6 +310,10 @@ , maxView , minViewWithKey , maxViewWithKey + , findMin + , findMax + , deleteFindMin + , deleteFindMax -- * Debugging , valid diff --git a/src/Data/Map/Strict/Internal.hs b/src/Data/Map/Strict/Internal.hs --- a/src/Data/Map/Strict/Internal.hs +++ b/src/Data/Map/Strict/Internal.hs @@ -5,8 +5,6 @@ #endif {-# OPTIONS_HADDOCK not-home #-} -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Map.Strict.Internal @@ -97,10 +95,12 @@ -- ** Delete\/Update , delete + , pop , adjust , adjustWithKey , update , updateWithKey + , upsert , updateLookupWithKey , alter , alterF @@ -141,6 +141,7 @@ , runWhenMissing -- *** @WhenMatched@ tactics + , dropMatched , zipWithMaybeMatched , zipWithMatched @@ -192,6 +193,7 @@ , mapKeys , mapKeysWith , mapKeysMonotonic + , mapAssocsMonotonic -- * Folds , foldr @@ -213,6 +215,9 @@ , keysSet , argSet , fromSet + , fromSetA + , fromSetMaybe + , fromSetMaybeA , fromArgSet -- ** Lists @@ -220,6 +225,7 @@ , fromList , fromListWith , fromListWithKey + , fromListUpsert -- ** Ordered lists , toAscList @@ -227,10 +233,12 @@ , fromAscList , fromAscListWith , fromAscListWithKey + , fromAscListUpsert , fromDistinctAscList , fromDescList , fromDescListWith , fromDescListWithKey + , fromDescListUpsert , fromDistinctDescList -- * Filter @@ -308,6 +316,7 @@ , dropMissing , filterMissing , filterAMissing + , dropMatched , merge , mergeA , ascLinkTop @@ -325,9 +334,6 @@ , argSet , assocs , atKeyImpl -#ifdef __GLASGOW_HASKELL__ - , atKeyPlain -#endif , balance , balanceL , balanceR @@ -336,6 +342,7 @@ , elems , empty , delete + , pop , deleteAt , deleteFindMax , deleteFindMin @@ -411,17 +418,16 @@ import Control.Applicative (Const (..), liftA3) import Data.Semigroup (Arg (..)) +import Data.Set.Internal (Set) import qualified Data.Set.Internal as Set import qualified Data.Map.Internal as L -import Utils.Containers.Internal.StrictPair +import Utils.Containers.Internal.Strict (StrictPair(..), toPair) #ifdef __GLASGOW_HASKELL__ import Data.Coerce #endif -#ifdef __GLASGOW_HASKELL__ import Data.Functor.Identity (Identity (..)) -#endif import qualified Data.Foldable as Foldable @@ -472,11 +478,7 @@ LT -> balanceL ky y (go kx x l) r GT -> balanceR ky y l (go kx x r) EQ -> Bin sz kx x l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} -#else -{-# INLINE insert #-} -#endif -- | \(O(\log n)\). Insert with a function, combining new value and old value. -- @'insertWith' f key value mp@ @@ -500,11 +502,7 @@ LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> let !y' = f x y in Bin sy kx y' l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWith #-} -#else -{-# INLINE insertWith #-} -#endif insertWithR :: Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a insertWithR = go @@ -516,11 +514,7 @@ LT -> balanceL ky y (go f kx x l) r GT -> balanceR ky y l (go f kx x r) EQ -> let !y' = f y x in Bin sy ky y' l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithR #-} -#else -{-# INLINE insertWithR #-} -#endif -- | \(O(\log n)\). Insert with a function, combining key, new value and old value. -- @'insertWithKey' f key value mp@ @@ -550,11 +544,7 @@ GT -> balanceR ky y l (go f kx x r) EQ -> let !x' = f kx x y in Bin sy kx x' l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKey #-} -#else -{-# INLINE insertWithKey #-} -#endif insertWithKeyR :: Ord k => (k -> a -> a -> a) -> k -> a -> Map k a -> Map k a insertWithKeyR = go @@ -569,11 +559,7 @@ GT -> balanceR ky y l (go f kx x r) EQ -> let !y' = f ky y x in Bin sy ky y' l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insertWithKeyR #-} -#else -{-# INLINE insertWithKeyR #-} -#endif -- | \(O(\log n)\). Combines insert operation with old value retrieval. -- The expression (@'insertLookupWithKey' f k x map@) @@ -608,11 +594,7 @@ in found :*: balanceR ky y l r' EQ -> let x' = f kx x y in x' `seq` (Just y :*: Bin sy kx x' l r) -#if __GLASGOW_HASKELL__ {-# INLINABLE insertLookupWithKey #-} -#else -{-# INLINE insertLookupWithKey #-} -#endif {-------------------------------------------------------------------- Deletion @@ -628,11 +610,7 @@ adjust :: Ord k => (a -> a) -> k -> Map k a -> Map k a adjust f = adjustWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ {-# INLINABLE adjust #-} -#else -{-# INLINE adjust #-} -#endif -- | \(O(\log n)\). Adjust a value at a specific key. When the key is not -- a member of the map, the original map is returned. @@ -653,11 +631,7 @@ GT -> Bin sx kx x l (go f k r) EQ -> Bin sx kx x' l r where !x' = f kx x -#if __GLASGOW_HASKELL__ {-# INLINABLE adjustWithKey #-} -#else -{-# INLINE adjustWithKey #-} -#endif -- | \(O(\log n)\). The expression (@'update' f k map@) updates the value @x@ -- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is @@ -670,11 +644,7 @@ update :: Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a update f = updateWithKey (\_ x -> f x) -#if __GLASGOW_HASKELL__ {-# INLINABLE update #-} -#else -{-# INLINE update #-} -#endif -- | \(O(\log n)\). The expression (@'updateWithKey' f k map@) updates the -- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', @@ -699,12 +669,27 @@ EQ -> case f kx x of Just x' -> x' `seq` Bin sx kx x' l r Nothing -> glue l r -#if __GLASGOW_HASKELL__ {-# INLINABLE updateWithKey #-} -#else -{-# INLINE updateWithKey #-} -#endif +-- | \(O(\log n)\). Update the value at a key or insert a value if the key is +-- not in the map. +-- +-- @ +-- let inc = maybe 1 (+1) +-- upsert inc \'a\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',2),(\'c\',2)] +-- upsert inc \'b\' (fromList [(\'a\',1),(\'c\',2)]) == fromList [(\'a\',1),(\'b\',1),(\'c\',2)] +-- @ +-- +-- @since 0.8.1 +upsert :: Ord k => (Maybe a -> a) -> k -> Map k a -> Map k a +upsert f !k (Bin sz kx x l r) = + case compare k kx of + LT -> balanceL kx x (upsert f k l) r + EQ -> let !x' = f (Just x) in Bin sz kx x' l r + GT -> balanceR kx x l (upsert f k r) +upsert f !k Tip = singleton k (f Nothing) +{-# INLINABLE upsert #-} + -- | \(O(\log n)\). Look up and update. See also 'updateWithKey'. -- This function returns the changed value, if it is updated. -- Returns the original key value if the map entry is deleted. @@ -729,11 +714,7 @@ EQ -> case f kx x of Just x' -> x' `seq` (Just x' :*: Bin sx kx x' l r) Nothing -> (Just x :*: glue l r) -#if __GLASGOW_HASKELL__ {-# INLINABLE updateLookupWithKey #-} -#else -{-# INLINE updateLookupWithKey #-} -#endif -- | \(O(\log n)\). The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. -- 'alter' can be used to insert, delete, or update a value in a 'Map'. @@ -764,31 +745,12 @@ EQ -> case f (Just x) of Just x' -> x' `seq` Bin sx kx x' l r Nothing -> glue l r -#if __GLASGOW_HASKELL__ {-# INLINABLE alter #-} -#else -{-# INLINE alter #-} -#endif -- | \(O(\log n)\). The expression (@'alterF' f k map@) alters the value @x@ at @k@, or absence thereof. -- 'alterF' can be used to inspect, insert, delete, or update a value in a 'Map'. -- In short: @'lookup' k \<$\> 'alterF' f k m = f ('lookup' k m)@. -- --- Example: --- --- @ --- interactiveAlter :: Int -> Map Int String -> IO (Map Int String) --- interactiveAlter k m = alterF f k m where --- f Nothing = do --- putStrLn $ show k ++ --- " was not found in the map. Would you like to add it?" --- getUserResponse1 :: IO (Maybe String) --- f (Just old) = do --- putStrLn $ "The key is currently bound to " ++ show old ++ --- ". Would you like to change or delete it?" --- getUserResponse2 :: IO (Maybe String) --- @ --- -- 'alterF' is the most general operation for working with an individual -- key that may or may not be in a given map. When used with trivial -- functors like 'Identity' and 'Const', it is often slightly slower than @@ -809,15 +771,27 @@ -- Note: 'alterF' is a flipped version of the @at@ combinator from -- @Control.Lens.At@. -- +-- === Examples +-- +-- @ +-- -- Lookup the value at the key, and also remove the existing value or set a new value. +-- lookupAndSet :: Ord k => k -> Maybe a -> Map k a -> (Maybe a, Map k a) +-- lookupAndSet k new = alterF (\\old -> (old, new)) k +-- @ +-- +-- @ +-- -- Delete the value at the key. If it is absent the result is Nothing. +-- mustDelete :: Ord k => k -> Map k a -> Maybe (Map k a) +-- mustDelete = alterF (Nothing <$) +-- @ +-- -- @since 0.5.8 alterF :: (Functor f, Ord k) => (Maybe a -> f (Maybe a)) -> k -> Map k a -> f (Map k a) alterF f k m = atKeyImpl Strict k f m -#ifndef __GLASGOW_HASKELL__ -{-# INLINE alterF #-} -#else -{-# INLINABLE [2] alterF #-} +#ifdef __GLASGOW_HASKELL__ +{-# INLINE [2] alterF #-} -- We can save a little time by recognizing the special case of -- `Control.Applicative.Const` and just doing a lookup. @@ -827,7 +801,7 @@ #-} atKeyIdentity :: Ord k => k -> (Maybe a -> Identity (Maybe a)) -> Map k a -> Identity (Map k a) -atKeyIdentity k f t = Identity $ atKeyPlain Strict k (coerce f) t +atKeyIdentity k f t = Identity (alter (coerce f) k t) {-# INLINABLE atKeyIdentity #-} #endif @@ -838,6 +812,8 @@ -- | \(O(\log n)\). Update the element at /index/. Calls 'error' when an -- invalid index is used. -- +-- __Note__: This function is partial. +-- -- > 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 @@ -920,9 +896,7 @@ unionsWith :: (Foldable f, Ord k) => (a->a->a) -> f (Map k a) -> Map k a unionsWith f ts = Foldable.foldl' (unionWith f) empty ts -#if __GLASGOW_HASKELL__ {-# INLINABLE unionsWith #-} -#endif {-------------------------------------------------------------------- Union with a combining function @@ -941,9 +915,7 @@ unionWith f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of (l2, mb, r2) -> link k1 x1' (unionWith f l1 l2) (unionWith f r1 r2) where !x1' = maybe x1 (f x1) mb -#if __GLASGOW_HASKELL__ {-# INLINABLE unionWith #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). -- Union with a combining function. @@ -961,9 +933,7 @@ unionWithKey f (Bin _ k1 x1 l1 r1) t2 = case splitLookup k1 t2 of (l2, mb, r2) -> link k1 x1' (unionWithKey f l1 l2) (unionWithKey f r1 r2) where !x1' = maybe x1 (f k1 x1) mb -#if __GLASGOW_HASKELL__ {-# INLINABLE unionWithKey #-} -#endif {-------------------------------------------------------------------- Difference @@ -981,9 +951,7 @@ differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a differenceWith f = merge preserveMissing dropMissing (zipWithMaybeMatched $ \_ x1 x2 -> f x1 x2) -#if __GLASGOW_HASKELL__ {-# INLINABLE differenceWith #-} -#endif -- | \(O(n+m)\). Difference with a combining function. When two equal keys are -- encountered, the combining function is applied to the key and both values. @@ -996,9 +964,7 @@ differenceWithKey :: Ord k => (k -> a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a differenceWithKey f = merge preserveMissing dropMissing (zipWithMaybeMatched f) -#if __GLASGOW_HASKELL__ {-# INLINABLE differenceWithKey #-} -#endif {-------------------------------------------------------------------- @@ -1019,9 +985,7 @@ !(l2, mb, r2) = splitLookup k t2 !l1l2 = intersectionWith f l1 l2 !r1r2 = intersectionWith f r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE intersectionWith #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). Intersection with a combining function. -- @@ -1038,9 +1002,7 @@ !(l2, mb, r2) = splitLookup k t2 !l1l2 = intersectionWithKey f l1 l2 !r1r2 = intersectionWithKey f r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE intersectionWithKey #-} -#endif -- | Map covariantly over a @'WhenMissing' f k x@. mapWhenMissing :: Functor f => (a -> b) -> WhenMissing f k x a -> WhenMissing f k x b @@ -1258,6 +1220,7 @@ combine !l' mx !r' = case mx of Nothing -> link2 l' r' Just !x' -> link kx x' l' r' +{-# INLINABLE traverseMaybeWithKey #-} -- | \(O(n)\). Map values and separate the 'Left' and 'Right' results. -- @@ -1418,24 +1381,92 @@ mapKeysWith :: Ord k2 => (a -> a -> a) -> (k1->k2) -> Map k1 a -> Map k2 a mapKeysWith c f m = finishB (foldlWithKey' (\b kx x -> insertWithB c (f kx) x b) emptyB m) -#if __GLASGOW_HASKELL__ {-# INLINABLE mapKeysWith #-} -#endif +-- | \(O(n)\). Map over keys and values with a function @f@ that is +-- monotonically strictly increasing in the keys. That is, for keys @kx@ and +-- @ky@ and values @x@ and @y@, if @kx@ < @ky@ then +-- @fst (f kx x)@ < @fst (f ky y)@. +-- +-- __Warning__: This function should be used only if @f@ is monotonically +-- strictly increasing in the key. This precondition is not checked. +-- +-- @since 0.8.1 +mapAssocsMonotonic :: (k1 -> a1 -> (k2, a2)) -> Map k1 a1 -> Map k2 a2 +mapAssocsMonotonic f = go + where + go Tip = Tip + go (Bin sz k1 x1 l r) = case f k1 x1 of + (k2, !x2) -> Bin sz k2 x2 (go l) (go r) +-- See Note [INLINABLE to expose unfoldings] in Data.IntMap.Internal +{-# INLINABLE mapAssocsMonotonic #-} + {-------------------------------------------------------------------- Conversions --------------------------------------------------------------------} --- | \(O(n)\). Build a map from a set of keys and a function which for each key --- computes its value. +-- | \(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) +fromSet :: (k -> a) -> Set k -> Map k a +#ifdef __GLASGOW_HASKELL__ +fromSet = + (coerce :: ((k -> Identity a) -> Set k -> Identity (Map k a)) + -> (k -> a) -> Set k -> Map k a) + fromSetA +#else +fromSet f = runIdentity . fromSetA (pure . f) +#endif +-- | \(O(n)\). Build a map from a 'Set' of keys and a function which for each +-- key computes its value in an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- > let f k = if k == 0 then Nothing else Just (6 `div` k) +-- > fromSetA f (Data.Set.fromList [1,2,3,4]) == Just (fromList [(1,6),(2,3),(3,2),(4,1)]) +-- > fromSetA f (Data.Set.fromList [0,1,2]) == Nothing +-- +-- @since 0.8.1 +fromSetA :: Applicative f => (k -> f a) -> Set k -> f (Map k a) +fromSetA _ Set.Tip = pure Tip +fromSetA f (Set.Bin sz x l r) = + liftA3 (flip (Bin sz x $!)) (fromSetA f l) (f x) (fromSetA f r) +{-# INLINABLE fromSetA #-} + +-- | \(O(n)\). Build a map from a 'Set' of keys and a function which for each +-- key optionally computes its value. +-- +-- > let f k = if even k then Just (replicate k 'a') else Nothing +-- > fromSetMaybe f (Data.Set.fromList [1,2,3,4]) == fromList [(2,"aa"), (4,"aaaa")] +-- +-- @since 0.8.1 +fromSetMaybe :: (k -> Maybe a) -> Set k -> Map k a +#ifdef __GLASGOW_HASKELL__ +fromSetMaybe = + (coerce :: ((k -> Identity (Maybe a)) -> Set k -> Identity (Map k a)) + -> (k -> Maybe a) -> Set k -> Map k a) + fromSetMaybeA +#else +fromSetMaybe f s = runIdentity (fromSetMaybeA (Identity . f) s) +#endif + +-- | \(O(n)\). Build a map from a 'Set' of keys and a function which for each +-- key optionally computes its values in an 'Applicative' context. +-- +-- The @Applicative@ actions are sequenced in order of increasing key. +-- +-- @since 0.8.1 +fromSetMaybeA :: Applicative f => (k -> f (Maybe a)) -> Set k -> f (Map k a) +fromSetMaybeA f = go + where + go Set.Tip = pure Tip + go (Set.Bin _ k l r) = + liftA3 (\l' mx r' -> maybe link2 (link k $!) mx l' r') (go l) (f k) (go r) +{-# INLINABLE fromSetMaybeA #-} + -- | \(O(n)\). Build a map from a set of elements contained inside 'Arg's. -- -- > fromArgSet (Data.Set.fromList [Arg 3 "aaa", Arg 5 "aaaaa"]) == fromList [(5,"aaaaa"), (3,"aaa")] @@ -1448,7 +1479,7 @@ {-------------------------------------------------------------------- Lists --------------------------------------------------------------------} --- | \(O(n \log n)\). Build a map from a list of key\/value pairs. See also 'fromAscList'. +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs. -- If the list contains more than one value for the same key, the last value -- for the key is retained. -- @@ -1463,7 +1494,7 @@ finishB (Foldable.foldl' (\b (kx, !x) -> insertB kx x b) emptyB xs) {-# INLINE fromList #-} -- INLINE for fusion --- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. -- -- If the keys are in non-decreasing order, this function takes \(O(n)\) time. -- @@ -1474,6 +1505,8 @@ -- -- The symmetric combining function @f@ is applied in a left-fold over the list, as @f new old@. -- +-- See also: 'fromListUpsert' +-- -- === Performance -- -- You should ensure that the given @f@ is fast with this order of arguments. @@ -1506,7 +1539,7 @@ finishB (Foldable.foldl' (\b (kx, x) -> insertWithB f kx x b) emptyB xs) {-# INLINE fromListWith #-} -- INLINE for fusion --- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. See also 'fromAscListWithKey'. +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a combining function. -- -- If the keys are in non-decreasing order, this function takes \(O(n)\) time. -- @@ -1515,12 +1548,35 @@ -- > fromListWithKey f [] == empty -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromListUpsert' fromListWithKey :: Ord k => (k -> a -> a -> a) -> [(k,a)] -> Map k a fromListWithKey f xs = finishB (Foldable.foldl' (\b (kx, x) -> insertWithB (f kx) kx x b) emptyB xs) {-# INLINE fromListWithKey #-} -- INLINE for fusion +-- | \(O(n \log n)\). Build a map from a list of key\/value pairs with a +-- combining function. +-- +-- If the keys are in non-decreasing order, this function takes \(O(n)\) time. +-- +-- The result is equivalent to performing an @upsert@ for every key\/value in +-- the list. +-- +-- @ +-- fromListUpsert f = foldl' (\\m (k, x) -> 'upsert' (f x) k m) 'empty' +-- @ +-- +-- > let f x = maybe [x] (x:) +-- > fromListUpsert f [(5,'a'), (5,'b'), (3,'c'), (3,'d'), (5,'e')] == fromList [(3,"dc"), (5,"eba")] +-- +-- @since 0.8.1 +fromListUpsert :: Ord k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b +fromListUpsert f xs = + finishB (Foldable.foldl' (\b (kx, x) -> upsertB (f x) kx b) emptyB xs) +{-# INLINE fromListUpsert #-} -- INLINE for fusion + {-------------------------------------------------------------------- Building trees from ascending/descending lists can be done in linear time. @@ -1574,6 +1630,8 @@ -- > valid (fromAscListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' fromAscListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a fromAscListWith f xs @@ -1591,6 +1649,8 @@ -- > valid (fromDescListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromDescListUpsert' fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a fromDescListWith f xs @@ -1605,11 +1665,13 @@ -- if the precondition may not hold. -- -- > 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 f [(3,"b"), (5,"a"), (5,"b"), (5,"c")] == fromList [(3, "b"), (5, "5:c5:ba")] +-- > valid (fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"c")]) == True +-- > valid (fromAscListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"c")]) == False -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromAscListUpsert' fromAscListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a fromAscListWithKey f xs = ascLinkAll (Foldable.foldl' next Nada xs) @@ -1636,6 +1698,8 @@ -- > valid (fromDescListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False -- -- Also see the performance note on 'fromListWith'. +-- +-- See also: 'fromDescListUpsert' fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a fromDescListWithKey f xs = descLinkAll (Foldable.foldl' next Nada xs) @@ -1649,6 +1713,54 @@ push kx !x = Push kx x {-# INLINE fromDescListWithKey #-} -- INLINE for fusion +-- | \(O(n)\). Build a map from an ascending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-decreasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromAscListUpsert f [(3,'a'), (3,'b'), (5,'c'), (5,'d'), (5,'e')] == fromList [(3,"ba"), (5,"edc")] +-- +-- @since 0.8.1 +fromAscListUpsert :: Eq k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b +fromAscListUpsert f xs = ascLinkAll (Foldable.foldl' next Nada xs) + where + next stk (!ky, y) = case stk of + Push kx x l stk' + | ky == kx -> push' ky (f y (Just x)) l stk' + | Tip <- l -> let !y' = f y Nothing + in ascLinkTop stk' 1 (singleton kx x) ky y' + | otherwise -> push' ky (f y Nothing) Tip stk + Nada -> push' ky (f y Nothing) Tip stk + push' kx !x = Push kx x +{-# INLINE fromAscListUpsert #-} -- INLINE for fusion + +-- | \(O(n)\). Build a map from a descending list in linear time with a +-- combining function for equal keys. +-- +-- __Warning__: This function should be used only if the keys are in +-- non-increasing order. This precondition is not checked. Use 'fromListUpsert' +-- if the precondition may not hold. +-- +-- > let f x = maybe [x] (x:) +-- > fromDescListUpsert f [(5,'a'), (5,'b'), (5,'c'), (3,'d'), (3,'e')] == fromList [(3,"ed"), (5,"cba")] +-- +-- @since 0.8.1 +fromDescListUpsert :: Eq k => (a -> Maybe b -> b) -> [(k, a)] -> Map k b +fromDescListUpsert f xs = descLinkAll (Foldable.foldl' next Nada xs) + where + next stk (!ky, y) = case stk of + Push kx x r stk' + | ky == kx -> push' ky (f y (Just x)) r stk' + | Tip <- r -> let !y' = f y Nothing + in descLinkTop ky y' 1 (singleton kx x) stk' + | otherwise -> push' ky (f y Nothing) Tip stk + Nada -> push' ky (f y Nothing) Tip stk + push' kx !x = Push kx x +{-# INLINE fromDescListUpsert #-} -- INLINE for fusion + -- | \(O(n)\). Build a map from an ascending list of distinct elements in linear time. -- -- __Warning__: This function should be used only if the keys are in @@ -1708,3 +1820,21 @@ where push' kx !x = Push kx x {-# INLINE insertWithB #-} + +-- Upsert a key-value. The given function is used to generate the value based +-- on the existing value for the key. Strict in the inserted value. +upsertB :: Ord k => (Maybe b -> b) -> k -> MapBuilder k b -> MapBuilder k b +upsertB f !ky b = case b of + BAsc stk -> case stk of + Push kx x l stk' -> case compare ky kx of + LT -> BMap (upsert f ky (ascLinkAll stk)) + EQ -> BAsc (push' ky (f (Just x)) l stk') + GT -> case l of + Tip -> let !y = f Nothing + in BAsc (ascLinkTop stk' 1 (singleton kx x) ky y) + Bin{} -> BAsc (push' ky (f Nothing) Tip stk) + Nada -> BAsc (push' ky (f Nothing) Tip Nada) + BMap m -> BMap (upsert f ky m) + where + push' kx !x = Push kx x +{-# INLINE upsertB #-} diff --git a/src/Data/Sequence.hs b/src/Data/Sequence.hs --- a/src/Data/Sequence.hs +++ b/src/Data/Sequence.hs @@ -172,6 +172,8 @@ viewl, -- :: Seq a -> ViewL a ViewR(..), viewr, -- :: Seq a -> ViewR a + -- ** List + toList, -- * Scans scanl, -- :: (a -> b -> a) -> a -> Seq b -> Seq a scanl1, -- :: (a -> a -> a) -> Seq a -> Seq a @@ -192,6 +194,7 @@ 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 + mapMaybe, -- :: (a -> Maybe b) -> Seq a -> Seq b -- * Sorting sort, -- :: Ord a => Seq a -> Seq a sortBy, -- :: (a -> a -> Ordering) -> Seq a -> Seq a @@ -207,10 +210,13 @@ adjust', -- :: (a -> a) -> Int -> Seq a -> Seq a update, -- :: Int -> a -> Seq a -> Seq a take, -- :: Int -> Seq a -> Seq a + takeR, -- :: Int -> Seq a -> Seq a drop, -- :: Int -> Seq a -> Seq a + dropR, -- :: Int -> Seq a -> Seq a insertAt, -- :: Int -> a -> Seq a -> Seq a deleteAt, -- :: Int -> Seq a -> Seq a splitAt, -- :: Int -> Seq a -> (Seq a, Seq a) + splitAtR, -- :: 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 diff --git a/src/Data/Sequence/Internal.hs b/src/Data/Sequence/Internal.hs --- a/src/Data/Sequence/Internal.hs +++ b/src/Data/Sequence/Internal.hs @@ -6,7 +6,6 @@ {-# LANGUAGE DeriveGeneric #-} {-# LANGUAGE DeriveLift #-} {-# LANGUAGE StandaloneDeriving #-} -{-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE InstanceSigs #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TemplateHaskellQuotes #-} @@ -18,10 +17,8 @@ {-# LANGUAGE PatternSynonyms #-} {-# LANGUAGE ViewPatterns #-} #endif -{-# LANGUAGE PatternGuards #-} {-# OPTIONS_HADDOCK not-home #-} -{-# OPTIONS_GHC -fno-warn-incomplete-uni-patterns #-} ----------------------------------------------------------------------------- -- | @@ -108,6 +105,8 @@ viewl, -- :: Seq a -> ViewL a ViewR(..), viewr, -- :: Seq a -> ViewR a + -- ** List + toList, -- * Scans scanl, -- :: (a -> b -> a) -> a -> Seq b -> Seq a scanl1, -- :: (a -> a -> a) -> Seq a -> Seq a @@ -128,6 +127,7 @@ 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 + mapMaybe, -- :: (a -> Maybe b) -> Seq a -> Seq b -- * Indexing lookup, -- :: Int -> Seq a -> Maybe a (!?), -- :: Seq a -> Int -> Maybe a @@ -136,10 +136,13 @@ adjust', -- :: (a -> a) -> Int -> Seq a -> Seq a update, -- :: Int -> a -> Seq a -> Seq a take, -- :: Int -> Seq a -> Seq a + takeR, -- :: Int -> Seq a -> Seq a drop, -- :: Int -> Seq a -> Seq a + dropR, -- :: Int -> Seq a -> Seq a insertAt, -- :: Int -> a -> Seq a -> Seq a deleteAt, -- :: Int -> Seq a -> Seq a splitAt, -- :: Int -> Seq a -> (Seq a, Seq a) + splitAtR, -- :: 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 @@ -197,7 +200,7 @@ import Data.Monoid (Monoid(..)) import Data.Functor (Functor(..)) import Utils.Containers.Internal.State (State(..), execState) -import Data.Foldable (foldr', toList) +import Data.Foldable (foldr') import qualified Data.Foldable as F import qualified Data.Semigroup as Semigroup @@ -213,9 +216,13 @@ import GHC.Exts (build) import Data.Data import Data.String (IsString(..)) +# if __GLASGOW_HASKELL__ >= 914 +import qualified Language.Haskell.TH.Lift as TH +# else import qualified Language.Haskell.TH.Syntax as TH -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () +# endif import GHC.Generics (Generic, Generic1) -- Array stuff, with GHC.Arr on GHC @@ -231,7 +238,7 @@ import Data.Functor.Identity (Identity(..)) -import Utils.Containers.Internal.StrictPair (StrictPair (..), toPair) +import Utils.Containers.Internal.Strict (StrictPair (..), toPair) import Control.Monad.Zip (MonadZip (..)) import Control.Monad.Fix (MonadFix (..), fix) @@ -330,7 +337,8 @@ #ifdef __GLASGOW_HASKELL__ -- | @since 0.6.6 instance TH.Lift a => TH.Lift (Seq a) where -# if MIN_VERSION_template_haskell(2,16,0) +-- template-haskell >= 2.16 +# if __GLASGOW_HASKELL__ >= 810 liftTyped t = [|| coerceFT z ||] # else lift t = [| coerceFT z |] @@ -424,9 +432,7 @@ {-# INLINE null #-} instance Traversable Seq where -#if __GLASGOW_HASKELL__ {-# INLINABLE traverse #-} -#endif traverse _ (Seq EmptyT) = pure (Seq EmptyT) traverse f' (Seq (Single (Elem x'))) = (\x'' -> Seq (Single (Elem x''))) <$> f' x' @@ -1001,7 +1007,9 @@ -- | @mempty@ = 'empty' instance Monoid (Seq a) where mempty = empty +#if !MIN_VERSION_base(4,11,0) mappend = (Semigroup.<>) +#endif -- | @(<>)@ = '(><)' -- @@ -1092,9 +1100,7 @@ foldMapNodeN :: Monoid m => (Node a -> m) -> Node (Node a) -> m foldMapNodeN f t = foldNode (<>) f t -#if __GLASGOW_HASKELL__ {-# INLINABLE foldMap #-} -#endif foldr _ z' EmptyT = z' foldr f' z' (Single x') = x' `f'` z' @@ -1758,6 +1764,10 @@ singleton x = Seq (Single (Elem x)) -- | \( O(\log n) \). @replicate n x@ is a sequence consisting of @n@ copies of @x@. +-- +-- Calls 'error' if @n < 0@. +-- +-- __Note__: This function is partial. replicate :: Int -> a -> Seq a replicate n x | n >= 0 = runIdentity (replicateA n (Identity x)) @@ -1766,6 +1776,10 @@ -- | 'replicateA' is an 'Applicative' version of 'replicate', and makes -- \( O(\log n) \) calls to 'liftA2' and 'pure'. -- +-- Calls 'error' if @n < 0@. +-- +-- __Note__: This function is partial. +-- -- > replicateA n x = sequenceA (replicate n x) replicateA :: Applicative f => Int -> f a -> f (Seq a) replicateA n x @@ -1773,13 +1787,9 @@ | otherwise = error "replicateA takes a nonnegative integer argument" {-# SPECIALIZE replicateA :: Int -> State a b -> State a (Seq b) #-} --- | 'replicateM' is the @Seq@ counterpart of --- @Control.Monad.'Control.Monad.replicateM'@. --- --- > replicateM n x = sequence (replicate n x) +-- | Synonym for 'replicateA'. -- --- For @base >= 4.8.0@ and @containers >= 0.5.11@, 'replicateM' --- is a synonym for 'replicateA'. +-- This definition exists for backwards compatibility. replicateM :: Applicative m => Int -> m a -> m (Seq a) replicateM = replicateA @@ -1788,11 +1798,15 @@ -- @k@ is 0. -- -- prop> cycleTaking k = fromList . take k . cycle . toList - +-- -- If you wish to concatenate a possibly empty sequence @xs@ with --- itself precisely @k@ times, use @'stimes' k xs@ instead of this --- function. +-- itself precisely @k@ times, use @'Data.Semigroup.stimes' k xs@ instead of +-- this function. -- +-- Calls 'error' if @k > 0@ and @null xs@. +-- +-- __Note__: This function is partial. +-- -- @since 0.5.8 cycleTaking :: Int -> Seq a -> Seq a cycleTaking n !_xs | n <= 0 = empty @@ -2205,7 +2219,9 @@ -- | \( 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)) +-- Calls 'error' if @n < 0@. +-- +-- __Note__: This function is partial. iterateN :: Int -> (a -> a) -> a -> Seq a iterateN n f x | n >= 0 = replicateA n (State (\ y -> (f y, y))) `execState` x @@ -2362,6 +2378,12 @@ viewRTree (Deep s pr m (Four w x y z)) = SnocRTree (Deep (s - size z) pr m (Three w x y)) z +-- | \(O(n)\). Convert to a list of elements. +-- +-- @since 0.8.1 +toList :: Seq a -> [a] +toList = F.toList + ------------------------------------------------------------------------ -- Scans -- @@ -2384,6 +2406,10 @@ -- | 'scanl1' is a variant of 'scanl' that has no starting value argument: -- +-- Calls 'error' if the sequence is empty. +-- +-- __Note__: This function is partial. +-- -- > 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 @@ -2395,6 +2421,10 @@ 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. +-- +-- Calls 'error' if the sequence is empty. +-- +-- __Note__: This function is partial. 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" @@ -2409,7 +2439,9 @@ -- -- prop> xs `index` i = toList xs !! i -- --- Caution: 'index' necessarily delays retrieving the requested +-- __Note__: This function is partial. Prefer 'lookup'. +-- +-- __Note__: 'index' necessarily delays retrieving the requested -- element until the result is forced. It can therefore lead to a space -- leak if the result is stored, unforced, in another structure. To retrieve -- an element immediately without forcing it, use 'lookup' or '(!?)'. @@ -3260,9 +3292,7 @@ foldMapWithIndexNodeN :: Monoid m => (Int -> Node a -> m) -> Int -> Node (Node a) -> m foldMapWithIndexNodeN f i t = foldWithIndexNode (<>) f i t -#if __GLASGOW_HASKELL__ {-# INLINABLE foldMapWithIndex #-} -#endif -- | 'traverseWithIndex' is a version of 'traverse' that also offers -- access to the index of each element. @@ -3343,11 +3373,7 @@ #ifdef __GLASGOW_HASKELL__ {-# INLINABLE [1] traverseWithIndex #-} -#else -{-# INLINE [1] traverseWithIndex #-} -#endif -#ifdef __GLASGOW_HASKELL__ {-# RULES "travWithIndex/mapWithIndex" forall f g xs . traverseWithIndex f (mapWithIndex g xs) = traverseWithIndex (\k a -> f k (g k a)) xs @@ -3380,6 +3406,10 @@ -- | \( O(n) \). Convert a given sequence length and a function representing that -- sequence into a sequence. -- +-- Calls 'error' if @n < 0@. +-- +-- __Note__: This function is partial. +-- -- @since 0.5.6.2 fromFunction :: Int -> (Int -> a) -> Seq a fromFunction len f | len < 0 = error "Data.Sequence.fromFunction called with negative len" @@ -3451,6 +3481,15 @@ | i <= 0 = empty | otherwise = xs +-- | \( O(\log(\min(i,n-i))) \). The last @i@ elements of a sequence. +-- If @i@ is negative, @'takeR' i s@ yields the empty sequence. +-- If the sequence contains fewer than @i@ elements, the whole sequence +-- is returned. +-- +-- @since 0.8.1 +takeR :: Int -> Seq a -> Seq a +takeR i xs = drop (length xs - i) xs + takeTreeE :: Int -> FingerTree (Elem a) -> FingerTree (Elem a) takeTreeE !_i EmptyT = EmptyT takeTreeE i t@(Single _) @@ -3613,6 +3652,15 @@ | i <= 0 = xs | otherwise = empty +-- | \( O(\log(\min(i,n-i))) \). Elements of a sequence before the last @i@. +-- If @i@ is negative, @'dropR' i s@ yields the whole sequence. +-- If the sequence contains fewer than @i@ elements, the empty sequence +-- is returned. +-- +-- @since 0.8.1 +dropR :: Int -> Seq a -> Seq a +dropR i xs = take (length xs - i) xs + -- We implement `drop` using a "take from the rear" strategy. There's no -- particular technical reason for this; it just lets us reuse the arithmetic -- from `take` (which itself reuses the arithmetic from `splitAt`) instead of @@ -3782,6 +3830,14 @@ | i <= 0 = (empty, xs) | otherwise = (xs, empty) +-- | \( O(\log(\min(i,n-i))) \). Split a sequence at a given position, +-- with the position being counted from the last (rightmost) element. +-- @'splitAtR' i s = ('dropR' i s, 'takeR' i s)@. +-- +-- @since 0.8.1 +splitAtR :: Int -> Seq a -> (Seq a, Seq a) +splitAtR i xs = splitAt (length xs - i) xs + -- | \( O(\log(\min(i,n-i))) \) A version of 'splitAt' that does not attempt to -- enhance sharing when the split point is less than or equal to 0, and that -- gives completely wrong answers when the split point is at least the length @@ -3958,6 +4014,10 @@ -- \( c = n \)) to \( O(n) \) (for \( c = 1 \)). The true bound is more like -- \( O \Bigl( \bigl(\frac{n}{c} - 1\bigr) (\log (c + 1)) + 1 \Bigr) \) -- +-- Calls 'error' if @n <= 0@ and @not (null xs)@. +-- +-- __Note__: This function is partial. +-- -- @since 0.5.8 chunksOf :: Int -> Seq a -> Seq (Seq a) chunksOf n xs | n <= 0 = @@ -4061,8 +4121,10 @@ (tailsTree f' m) (fmap (f . digitToTree) (tailsDigit sf)) where - f' ms = let ConsLTree node m' = viewLTree ms in + f' ms = case viewLTree ms of + ConsLTree node m' -> fmap (\ pr' -> f (deep pr' m' sf)) (tailsNode node) + EmptyLTree -> error "EmptyLTree" {-# 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) #-} @@ -4076,8 +4138,10 @@ (initsTree f' m) (fmap (f . deep pr m) (initsDigit sf)) where - f' ms = let SnocRTree m' node = viewRTree ms in + f' ms = case viewRTree ms of + SnocRTree m' node -> fmap (\ sf' -> f (deep pr m' sf')) (initsNode node) + EmptyRTree -> error "EmptyRTree" {-# INLINE foldlWithIndex #-} -- | 'foldlWithIndex' is a version of 'foldl' that also provides access @@ -4169,6 +4233,15 @@ filter :: (a -> Bool) -> Seq a -> Seq a filter p = foldl' (\ xs x -> if p x then xs `snoc'` x else xs) empty +-- | \( O(n) \). Map elements and collect the 'Just' results. +-- +-- @since 0.8.1 +mapMaybe :: (a -> Maybe b) -> Seq a -> Seq b +mapMaybe f = foldl' go empty + where go xs x = case f x of + Nothing -> xs + Just x' -> xs `snoc'` x' + -- Indexing sequences -- | 'elemIndexL' finds the leftmost index of the specified element, @@ -4278,8 +4351,10 @@ -- eventually find a less mind-bending way to accomplish this. -- | \( O(n) \). Create a sequence from a finite list of elements. --- There is a function 'toList' in the opposite direction for all --- instances of the 'Foldable' class, including 'Seq'. +-- +-- @'fromList' . 'toList' = id@ +-- +-- For any finite list @xs@, @'toList' ('fromList' xs) = xs@. fromList :: [a] -> Seq a -- Note: we can avoid map_elem if we wish by scattering -- Elem applications throughout mkTreeE and getNodesE, but diff --git a/src/Data/Sequence/Internal/Sorting.hs b/src/Data/Sequence/Internal/Sorting.hs --- a/src/Data/Sequence/Internal/Sorting.hs +++ b/src/Data/Sequence/Internal/Sorting.hs @@ -407,8 +407,8 @@ -> Maybe b foldToMaybeWithIndexTree = foldToMaybeWithIndexTree' where - {-# SPECIALISE foldToMaybeWithIndexTree' :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> FingerTree (Elem y) -> Maybe b #-} - {-# SPECIALISE foldToMaybeWithIndexTree' :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> FingerTree (Node y) -> Maybe b #-} + {-# SPECIALIZE foldToMaybeWithIndexTree' :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> FingerTree (Elem y) -> Maybe b #-} + {-# SPECIALIZE foldToMaybeWithIndexTree' :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> FingerTree (Node y) -> Maybe b #-} foldToMaybeWithIndexTree' :: Sized a => (b -> b -> b) -> (Int -> a -> b) -> Int -> FingerTree a -> Maybe b @@ -422,14 +422,14 @@ m' = foldToMaybeWithIndexTree' (<+>) (node (<+>) f) sPspr m !sPspr = s + size pr !sPsprm = sPspr + size m - {-# SPECIALISE digit :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> Digit (Elem y) -> b #-} - {-# SPECIALISE digit :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> Digit (Node y) -> b #-} + {-# SPECIALIZE digit :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> Digit (Elem y) -> b #-} + {-# SPECIALIZE digit :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> Digit (Node y) -> b #-} digit :: Sized a => (b -> b -> b) -> (Int -> a -> b) -> Int -> Digit a -> b digit = foldWithIndexDigit - {-# SPECIALISE node :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> Node (Elem y) -> b #-} - {-# SPECIALISE node :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> Node (Node y) -> b #-} + {-# SPECIALIZE node :: (b -> b -> b) -> (Int -> Elem y -> b) -> Int -> Node (Elem y) -> b #-} + {-# SPECIALIZE node :: (b -> b -> b) -> (Int -> Node y -> b) -> Int -> Node (Node y) -> b #-} node :: Sized a => (b -> b -> b) -> (Int -> a -> b) -> Int -> Node a -> b diff --git a/src/Data/Set.hs b/src/Data/Set.hs --- a/src/Data/Set.hs +++ b/src/Data/Set.hs @@ -3,8 +3,6 @@ {-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Set @@ -17,9 +15,19 @@ -- = Finite Sets -- -- The @'Set' e@ type represents a set of elements of type @e@. Most operations --- require that @e@ be an instance of the 'Ord' class. A 'Set' is strict in its --- elements. +-- require that @e@ have an instance of the 'Ord' class. A 'Set' is strict in +-- its elements. -- +-- When deciding if this is the correct data structure to use, consider: +-- +-- * If you are using 'Int' keys, you will get much better performance for most +-- operations using "Data.IntSet". +-- +-- * If you don't care about ordering and don't handle untrusted elements, +-- consider using @Data.HashSet@ from the +-- +-- package instead. +-- -- For a walkthrough of the most commonly used functions see the -- . -- @@ -29,11 +37,11 @@ -- > import Data.Set (Set) -- > import qualified Data.Set as Set -- --- Note that the implementation is generally /left-biased/. Functions that take --- two sets as arguments and combine them, such as `union` and `intersection`, --- prefer the entries in the first argument to those in the second. Of course, --- this bias can only be observed when equality is an equivalence relation --- instead of structural equality. +-- The @'Ord' e@ instance is expected to be lawful and define a total order. +-- Unless otherwise specified, operations expect equality to be extensional: if +-- elements @x1@ and @x2@ satisfy @x1 == x2@, they are considered identical. For +-- instance, if only one element must be retained by an operation, it is free to +-- select either. -- -- -- == Warning @@ -102,6 +110,7 @@ -- * Deletion , delete + , pop -- * Generalized insertion/deletion @@ -134,9 +143,11 @@ -- * Filter , S.filter + , filterA , takeWhileAntitone , dropWhileAntitone , spanAntitone + , mapMaybe , partition , split , splitMember @@ -167,14 +178,14 @@ -- * Min\/Max , lookupMin , lookupMax - , findMin - , findMax , deleteMin , deleteMax - , deleteFindMin - , deleteFindMax , maxView , minView + , findMin + , findMax + , deleteFindMin + , deleteFindMax -- * Conversion diff --git a/src/Data/Set/Internal.hs b/src/Data/Set/Internal.hs --- a/src/Data/Set/Internal.hs +++ b/src/Data/Set/Internal.hs @@ -1,6 +1,5 @@ {-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} #ifdef __GLASGOW_HASKELL__ {-# LANGUAGE Trustworthy #-} {-# LANGUAGE DeriveLift #-} @@ -11,8 +10,6 @@ {-# OPTIONS_HADDOCK not-home #-} -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Set.Internal @@ -78,14 +75,11 @@ -- INLINABLE (that exposes the unfolding). --- [Note: Using INLINE] +-- 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. +-- We mark some functions INLINE where it is more beneficial than INLINABLE. +-- There is a price to pay -- code growth. The code INLINED is therefore usually +-- only the tree navigation, other work such as rebalancing is not INLINED. -- -- All methods marked INLINE have to be nonrecursive -- a 'go' function doing -- the real work is provided. @@ -142,6 +136,7 @@ , singleton , insert , delete + , pop , alterF , powerSet @@ -159,9 +154,11 @@ -- * Filter , filter + , filterA , takeWhileAntitone , dropWhileAntitone , spanAntitone + , mapMaybe , partition , split , splitMember @@ -221,26 +218,51 @@ , showTreeWith , valid + -- * General merge functions + , WhenMissing(..) + , SimpleWhenMissing + , dropMissing + , preserveMissing + , filterMissing + , filterAMissing + , whenMissing + , runWhenMissing + , WhenMatched(..) + , SimpleWhenMatched + , dropMatched + , preserveMatched + , filterMatched + , filterAMatched + , runWhenMatched + , merge + , mergeA + -- Internals (for testing) , bin , balanced , link - , merge + , link2 + , balanceL + , balanceR + , glue + , insertMin + , insertMax ) where import Utils.Containers.Internal.Prelude hiding (filter,foldl,foldl',foldr,null,map,take,drop,splitAt) import Prelude () -import Control.Applicative (Const(..)) +import Control.Applicative (Const(..), liftA3) import qualified Data.List as List import Data.Semigroup (Semigroup(..), stimesIdempotentMonoid, stimesIdempotent) import Data.Functor.Classes -import Data.Functor.Identity (Identity) +import Data.Functor.Identity (Identity(..)) import qualified Data.Foldable as Foldable import Control.DeepSeq (NFData(rnf),NFData1(liftRnf)) import Data.List.NonEmpty (NonEmpty(..)) -import Utils.Containers.Internal.StrictPair +import Utils.Containers.Internal.Strict + (StrictPair(..), StrictTriple(..), toPair) import Utils.Containers.Internal.PtrEquality import Utils.Containers.Internal.EqOrdUtil (EqM(..), OrdM(..)) @@ -252,9 +274,13 @@ import GHC.Exts ( build, lazy ) import qualified GHC.Exts as GHCExts import Data.Data +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else import Language.Haskell.TH.Syntax (Lift) -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () +# endif import Data.Coerce (coerce) #endif @@ -267,9 +293,7 @@ -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). See 'difference'. (\\) :: Ord a => Set a -> Set a -> Set a m1 \\ m2 = difference m1 m2 -#if __GLASGOW_HASKELL__ {-# INLINABLE (\\) #-} -#endif {-------------------------------------------------------------------- Sets are size balanced trees @@ -293,7 +317,9 @@ instance Ord a => Monoid (Set a) where mempty = empty mconcat = unions +#if !MIN_VERSION_base(4,11,0) mappend = (<>) +#endif -- | @(<>)@ = 'union' -- @@ -391,20 +417,12 @@ LT -> go x l GT -> go x r EQ -> True -#if __GLASGOW_HASKELL__ {-# INLINABLE member #-} -#else -{-# INLINE member #-} -#endif -- | \(O(\log n)\). Is the element not in the set? notMember :: Ord a => a -> Set a -> Bool notMember a t = not $ member a t -#if __GLASGOW_HASKELL__ {-# INLINABLE notMember #-} -#else -{-# INLINE notMember #-} -#endif -- | \(O(\log n)\). Find largest element smaller than the given one. -- @@ -420,11 +438,7 @@ 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__ {-# INLINABLE lookupLT #-} -#else -{-# INLINE lookupLT #-} -#endif -- | \(O(\log n)\). Find smallest element greater than the given one. -- @@ -440,11 +454,7 @@ 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__ {-# INLINABLE lookupGT #-} -#else -{-# INLINE lookupGT #-} -#endif -- | \(O(\log n)\). Find largest element smaller or equal to the given one. -- @@ -463,11 +473,7 @@ 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__ {-# INLINABLE lookupLE #-} -#else -{-# INLINE lookupLE #-} -#endif -- | \(O(\log n)\). Find smallest element greater or equal to the given one. -- @@ -486,11 +492,7 @@ 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__ {-# INLINABLE lookupGE #-} -#else -{-# INLINE lookupGE #-} -#endif {-------------------------------------------------------------------- Construction @@ -528,11 +530,7 @@ where !r' = go orig x r EQ | lazy orig `seq` (orig `ptrEq` y) -> t | otherwise -> Bin sz (lazy orig) l r -#if __GLASGOW_HASKELL__ {-# INLINABLE insert #-} -#else -{-# INLINE insert #-} -#endif #ifndef __GLASGOW_HASKELL__ lazy :: a -> a @@ -557,11 +555,7 @@ | otherwise -> balanceR y l r' where !r' = go orig x r EQ -> t -#if __GLASGOW_HASKELL__ {-# INLINABLE insertR #-} -#else -{-# INLINE insertR #-} -#endif -- | \(O(\log n)\). Delete an element from a set. @@ -579,12 +573,37 @@ | otherwise -> balanceL y l r' where !r' = go x r EQ -> glue l r -#if __GLASGOW_HASKELL__ {-# INLINABLE delete #-} -#else -{-# INLINE delete #-} -#endif +-- | \(O(\log n)\). Pop an element from the set. +-- +-- Returns @Nothing@ if the element is not a member of the set. Otherwise +-- returns @Just@ the set with the element removed. +-- +-- @ +-- pop 1 (fromList [0,2,4]) == Nothing +-- pop 2 (fromList [0,2,4]) == Just (fromList [0,4]) +-- @ +-- +-- @since 0.8.1 +pop :: Ord a => a -> Set a -> Maybe (Set a) +pop x0 t0 = case go x0 t0 of + True :*: t -> Just t + _ -> Nothing + where + -- We use `StrictPair Bool (Set a)` instead of a sum to avoid allocations. + -- See Note [Popped impl] in Data.Map.Internal + go !x (Bin _ y l r) = case compare x y of + LT -> case go x l of + True :*: l' -> True :*: balanceR y l' r + q -> q + EQ -> True :*: glue l r + GT -> case go x r of + True :*: r' -> True :*: balanceL y l r' + q -> q + go !_ Tip = False :*: Tip +{-# INLINABLE pop #-} + -- | \(O(\log n)\) @('alterF' f x s)@ can delete or insert @x@ in @s@ depending on -- whether an equal element is found in @s@. -- @@ -599,55 +618,62 @@ -- -- Note: 'alterF' is a variant of the @at@ combinator from "Control.Lens.At". -- +-- === Examples +-- +-- @ +-- -- Get whether the element is a member, and also insert or remove it. +-- getAndSet :: Ord a => a -> Bool -> Set a -> (Bool, Set a) +-- getAndSet x new = alterF (\\old -> (old, new)) x +-- @ +-- +-- @ +-- -- Delete the element. If it is absent the result is Nothing. +-- mustDelete :: Ord a => a -> Set a -> Maybe (Set a) +-- mustDelete = alterF (\\b -> if b then Just False else Nothing) +-- @ +-- -- @since 0.6.3.1 + +-- See Note [alterF implementation] alterF :: (Ord a, Functor f) => (Bool -> f Bool) -> a -> Set a -> f (Set a) alterF f k s = fmap choose (f member_) where - (member_, inserted, deleted) = case alteredSet k s of - Deleted d -> (True , s, d) - Inserted i -> (False, i, s) - + MemberIndex member_ i = memberIndex k s + inserted = if member_ then s else insertAt i k s + deleted = if member_ then deleteAt i s else s choose True = inserted choose False = deleted -#ifndef __GLASGOW_HASKELL__ -{-# INLINE alterF #-} -#else -{-# INLINABLE [2] alterF #-} +#ifdef __GLASGOW_HASKELL__ +{-# INLINE [2] alterF #-} {-# RULES "alterF/Const" forall k (f :: Bool -> Const a Bool) . alterF f k = \s -> Const . getConst . f $ member k s #-} #endif -{-# SPECIALIZE alterF :: Ord a => (Bool -> Identity Bool) -> a -> Set a -> Identity (Set a) #-} - -data AlteredSet a - -- | The needle is present in the original set. - -- We return the set where the needle is deleted. - = Deleted !(Set a) +data MemberIndex = MemberIndex !Bool {-# UNPACK #-} !Int - -- | The needle is not present in the original set. - -- We return the set with the needle inserted. - | Inserted !(Set a) +-- Whether the element is a member of the set along with its index. +-- If it is not a member, the index is the index it would have if inserted. +memberIndex :: Ord a => a -> Set a -> MemberIndex +memberIndex = go 0 + where + go !i !_ Tip = MemberIndex False i + go !i !x (Bin _ y l r) = case compare x y of + LT -> go i x l + EQ -> MemberIndex True (i + size l) + GT -> go (i + size l + 1) x r +{-# INLINABLE memberIndex #-} -alteredSet :: Ord a => a -> Set a -> AlteredSet a -alteredSet x0 s0 = go x0 s0 +-- Insert the element at the given index. The caller must ensure that the index +-- is correct and will not violate Set invariants. +insertAt :: Int -> a -> Set a -> Set a +insertAt !_ !x Tip = singleton x +insertAt !i !x (Bin _ y l r) + | i <= sizeL = balanceL y (insertAt i x l) r + | otherwise = balanceR y l (insertAt (i - sizeL - 1) x r) where - go :: Ord a => a -> Set a -> AlteredSet a - go x Tip = Inserted (singleton x) - go x (Bin _ y l r) = case compare x y of - LT -> case go x l of - Deleted d -> Deleted (balanceR y d r) - Inserted i -> Inserted (balanceL y i r) - GT -> case go x r of - Deleted d -> Deleted (balanceL y l d) - Inserted i -> Inserted (balanceR y l i) - EQ -> Deleted (glue l r) -#if __GLASGOW_HASKELL__ -{-# INLINABLE alteredSet #-} -#else -{-# INLINE alteredSet #-} -#endif + !sizeL = size l {-------------------------------------------------------------------- Subset @@ -662,9 +688,7 @@ isProperSubsetOf :: Ord a => Set a -> Set a -> Bool isProperSubsetOf s1 s2 = size s1 < size s2 && isSubsetOfX s1 s2 -#if __GLASGOW_HASKELL__ {-# INLINABLE isProperSubsetOf #-} -#endif -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). @@ -679,9 +703,7 @@ isSubsetOf :: Ord a => Set a -> Set a -> Bool isSubsetOf t1 t2 = size t1 <= size t2 && isSubsetOfX t1 t2 -#if __GLASGOW_HASKELL__ {-# INLINABLE isSubsetOf #-} -#endif -- Test whether a set is a subset of another without the *initial* -- size test. @@ -715,9 +737,7 @@ isSubsetOfX l lt && isSubsetOfX r gt where (lt,found,gt) = splitMember x t -#if __GLASGOW_HASKELL__ {-# INLINABLE isSubsetOfX #-} -#endif {-------------------------------------------------------------------- Disjoint @@ -774,6 +794,8 @@ -- | \(O(\log n)\). The minimal element of the set. Calls 'error' if the set is -- empty. +-- +-- __Note__: This function is partial. Prefer 'lookupMin'. findMin :: Set a -> a findMin t | Just r <- lookupMin t = r @@ -795,6 +817,8 @@ -- | \(O(\log n)\). The maximal element of the set. Calls 'error' if the set is -- empty. +-- +-- __Note__: This function is partial. Prefer 'lookupMax'. findMax :: Set a -> a findMax t | Just r <- lookupMax t = r @@ -818,9 +842,7 @@ -- | The union of the sets in a Foldable structure : (@'unions' == 'foldl' 'union' 'empty'@). unions :: (Foldable f, Ord a) => f (Set a) -> Set a unions = Foldable.foldl' union empty -#if __GLASGOW_HASKELL__ -{-# INLINABLE unions #-} -#endif +{-# INLINE unions #-} -- Inline for list fusion -- | \(O\bigl(m \log\bigl(\frac{n}{m}+1\bigr)\bigr), \; 0 < m \leq n\). The union of two sets, preferring the first set when -- equal elements are encountered. @@ -835,9 +857,7 @@ | otherwise -> link x l1l2 r1r2 where !l1l2 = union l1 l2 !r1r2 = union r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE union #-} -#endif {-------------------------------------------------------------------- Difference @@ -853,12 +873,10 @@ difference t1 (Bin _ x l2 r2) = case split x t1 of (l1, r1) | size l1l2 + size r1r2 == size t1 -> t1 - | otherwise -> merge l1l2 r1r2 + | otherwise -> link2 l1l2 r1r2 where !l1l2 = difference l1 l2 !r1r2 = difference r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE difference #-} -#endif {-------------------------------------------------------------------- Intersection @@ -881,14 +899,12 @@ | b = if l1l2 `ptrEq` l1 && r1r2 `ptrEq` r1 then t1 else link x l1l2 r1r2 - | otherwise = merge l1l2 r1r2 + | otherwise = link2 l1l2 r1r2 where !(l2, b, r2) = splitMember x t2 !l1l2 = intersection l1 l2 !r1r2 = intersection r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE intersection #-} -#endif -- | The intersection of a series of sets. Intersections are performed -- left-to-right. @@ -945,31 +961,51 @@ symmetricDifference Tip t2 = t2 symmetricDifference t1 Tip = t1 symmetricDifference (Bin _ x l1 r1) t2 - | found = merge l1l2 r1r2 + | found = link2 l1l2 r1r2 | otherwise = link x l1l2 r1r2 where !(l2, found, r2) = splitMember x t2 !l1l2 = symmetricDifference l1 l2 !r1r2 = symmetricDifference r1 r2 -#if __GLASGOW_HASKELL__ {-# INLINABLE symmetricDifference #-} -#endif {-------------------------------------------------------------------- Filter and partition --------------------------------------------------------------------} --- | \(O(n)\). Filter all elements that satisfy the predicate. +-- | \(O(n)\). Keep all elements that satisfy the predicate. filter :: (a -> Bool) -> Set a -> Set a filter _ Tip = Tip filter p t@(Bin _ x l r) | p x = if l `ptrEq` l' && r `ptrEq` r' then t else link x l' r' - | otherwise = merge l' r' + | otherwise = link2 l' r' where !l' = filter p l !r' = filter p r +-- | \(O(n)\). Keep all elements that satisfy the Applicative predicate. +-- +-- @since 0.8.1 +filterA :: Applicative f => (a -> f Bool) -> Set a -> f (Set a) +filterA p = go + where + -- Handle the singleton case separately as it is possible for fmap to be + -- cheaper than liftA3. + go t@(Bin 1 x _ _) = finish <$> p x + where + finish False = Tip + finish True = t + go t@(Bin _ x l r) = liftA3 doLink (go l) (p x) (go r) + where + doLink !l' keepx !r' + | keepx = if l `ptrEq` l' && r `ptrEq` r' + then t + else link x l' r' + | otherwise = link2 l' r' + go Tip = pure Tip +{-# INLINABLE filterA #-} + -- | \(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'. @@ -981,12 +1017,29 @@ ((l1 :*: l2), (r1 :*: r2)) | p x -> (if l1 `ptrEq` l && r1 `ptrEq` r then t - else link x l1 r1) :*: merge l2 r2 - | otherwise -> merge l1 r1 :*: + else link x l1 r1) :*: link2 l2 r2 + | otherwise -> link2 l1 r1 :*: (if l2 `ptrEq` l && r2 `ptrEq` r then t else link x l2 r2) +{-------------------------------------------------------------------- + Maybes +--------------------------------------------------------------------} + +-- | \(O(n \log n)\). Map elements and collect the 'Just' results. +-- +-- If the function is monotonically non-decreasing, this function takes \(O(n)\) +-- time. +-- +-- @since 0.8.1 +mapMaybe :: Ord b => (a -> Maybe b) -> Set a -> Set b +mapMaybe f t = finishB (foldl' go emptyB t) + where go b x = case f x of + Nothing -> b + Just x' -> insertB x' b +{-# INLINABLE mapMaybe #-} + {---------------------------------------------------------------------- Map ----------------------------------------------------------------------} @@ -1001,9 +1054,7 @@ map :: Ord b => (a->b) -> Set a -> Set b map f t = finishB (foldl' (\b x -> insertB (f x) b) emptyB t) -#if __GLASGOW_HASKELL__ {-# INLINABLE map #-} -#endif -- | \(O(n)\). -- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly increasing. @@ -1213,7 +1264,7 @@ ascLinkTop stk !_ r y = Push y r stk ascLinkAll :: Stack a -> Set a -ascLinkAll stk = foldl'Stack (\r x l -> link x l r) Tip stk +ascLinkAll stk = foldl'Stack (\r x l -> linkL x l r) Tip stk {-# INLINABLE ascLinkAll #-} -- | \(O(n)\). Build a set from a descending list of distinct elements in linear time. @@ -1241,7 +1292,7 @@ descLinkTop y !_ r stk = Push y r stk descLinkAll :: Stack a -> Set a -descLinkAll stk = foldl'Stack (\l x r -> link x l r) Tip stk +descLinkAll stk = foldl'Stack (\l x r -> linkR x l r) Tip stk {-# INLINABLE descLinkAll #-} data Stack a = Push !a !(Set a) !(Stack a) | Nada @@ -1398,27 +1449,26 @@ splitS _ Tip = (Tip :*: Tip) splitS x (Bin _ y l r) = case compare x y of - LT -> let (lt :*: gt) = splitS x l in (lt :*: link y gt r) - GT -> let (lt :*: gt) = splitS x r in (link y l lt :*: gt) + LT -> let (lt :*: gt) = splitS x l in (lt :*: linkR y gt r) + GT -> let (lt :*: gt) = splitS x r in (linkL y l lt :*: gt) EQ -> (l :*: r) {-# INLINABLE splitS #-} -- | \(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 (lt, found, gt') - GT -> let (lt, found, gt) = splitMember x r - !lt' = link y l lt - in (lt', found, gt) - EQ -> (l, True, r) -#if __GLASGOW_HASKELL__ +splitMember x0 t = case go x0 t of + TripleS lt found gt -> (lt, found, gt) + where + go _ Tip = TripleS Tip False Tip + go x (Bin _ y l r) = + case compare x y of + LT -> case go x l of + TripleS lt found gt -> TripleS lt found (linkR y gt r) + GT -> case go x r of + TripleS lt found gt -> TripleS (linkL y l lt) found gt + EQ -> TripleS l True r {-# INLINABLE splitMember #-} -#endif {-------------------------------------------------------------------- Indexing @@ -1429,6 +1479,8 @@ -- to, but not including, the 'size' of the set. Calls 'error' when the element -- is not a 'member' of the set. -- +-- __Note__: This function is partial. Prefer 'lookupIndex'. +-- -- > findIndex 2 (fromList [5,3]) Error: element is not in the set -- > findIndex 3 (fromList [5,3]) == 0 -- > findIndex 5 (fromList [5,3]) == 1 @@ -1446,9 +1498,7 @@ LT -> go idx x l GT -> go (idx + size l + 1) x r EQ -> idx + size l -#if __GLASGOW_HASKELL__ {-# INLINABLE findIndex #-} -#endif -- | \(O(\log n)\). Look up 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 @@ -1471,14 +1521,14 @@ LT -> go idx x l GT -> go (idx + size l + 1) x r EQ -> Just $! idx + size l -#if __GLASGOW_HASKELL__ {-# INLINABLE lookupIndex #-} -#endif -- | \(O(\log n)\). Retrieve an element by its /index/, i.e. by its zero-based -- index in the 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. -- +-- __Note__: This function is partial. +-- -- > elemAt 0 (fromList [5,3]) == 3 -- > elemAt 1 (fromList [5,3]) == 5 -- > elemAt 2 (fromList [5,3]) Error: index out of range @@ -1499,6 +1549,8 @@ -- 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. -- +-- __Note__: This function is partial. +-- -- > deleteAt 0 (fromList [5,3]) == singleton 5 -- > deleteAt 1 (fromList [5,3]) == singleton 3 -- > deleteAt 2 (fromList [5,3]) Error: index out of range @@ -1534,7 +1586,7 @@ go i (Bin _ x l r) = case compare i sizeL of LT -> go i l - GT -> link x l (go (i - sizeL - 1) r) + GT -> linkL x l (go (i - sizeL - 1) r) EQ -> l where sizeL = size l @@ -1554,7 +1606,7 @@ go !_ Tip = Tip go i (Bin _ x l r) = case compare i sizeL of - LT -> link x (go i l) r + LT -> linkR x (go i l) r GT -> go (i - sizeL - 1) r EQ -> insertMin x r where sizeL = size l @@ -1574,9 +1626,9 @@ go i (Bin _ x l r) = case compare i sizeL of LT -> case go i l of - ll :*: lr -> ll :*: link x lr r + ll :*: lr -> ll :*: linkR x lr r GT -> case go (i - sizeL - 1) r of - rl :*: rr -> link x l rl :*: rr + rl :*: rr -> linkL x l rl :*: rr EQ -> l :*: insertMin x r where sizeL = size l @@ -1594,7 +1646,7 @@ takeWhileAntitone :: (a -> Bool) -> Set a -> Set a takeWhileAntitone _ Tip = Tip takeWhileAntitone p (Bin _ x l r) - | p x = link x l (takeWhileAntitone p r) + | p x = linkL x l (takeWhileAntitone p r) | otherwise = takeWhileAntitone p l -- | \(O(\log n)\). Drop while a predicate on the elements holds. @@ -1612,7 +1664,7 @@ dropWhileAntitone _ Tip = Tip dropWhileAntitone p (Bin _ x l r) | p x = dropWhileAntitone p r - | otherwise = link x (dropWhileAntitone p l) r + | otherwise = linkR x (dropWhileAntitone p l) r -- | \(O(\log n)\). Divide a set at the point where a predicate on the elements stops holding. -- The user is responsible for ensuring that for all elements @j@ and @k@ in the set, @@ -1635,8 +1687,8 @@ where go _ Tip = Tip :*: Tip go p (Bin _ x l r) - | p x = let u :*: v = go p r in link x l u :*: v - | otherwise = let u :*: v = go p l in u :*: link x v r + | p x = let u :*: v = go p r in linkL x l u :*: v + | otherwise = let u :*: v = go p l in u :*: linkR x v r {-------------------------------------------------------------------- SetBuilder @@ -1707,7 +1759,7 @@ 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. + [link2 l r] Merges two trees and restores balance. --------------------------------------------------------------------} {-------------------------------------------------------------------- @@ -1716,20 +1768,50 @@ 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 +link x l@(Bin lsz lx ll lr) r@(Bin rsz rx rl rr) + | delta*lsz < rsz = balanceL rx (linkR_ x lsz l rl) rr + | delta*rsz < lsz = balanceR lx ll (linkL_ x lr rsz r) + | otherwise = Bin (1+lsz+rsz) x l r +-- Variant of link. Restores balance when the left tree may be too large for the +-- right tree, but not the other way around. +linkL :: a -> Set a -> Set a -> Set a +linkL x l r = case r of + Tip -> insertMax x l + Bin rsz _ _ _ -> linkL_ x l rsz r +linkL_ :: a -> Set a -> Int -> Set a -> Set a +linkL_ x l !rsz r = case l of + Bin lsz lx ll lr + | delta*rsz < lsz -> balanceR lx ll (linkL_ x lr rsz r) + | otherwise -> Bin (1+lsz+rsz) x l r + Tip -> Bin (1+rsz) x Tip r + +-- Variant of link. Restores balance when the right tree may be too large for +-- the left tree, but not the other way around. +linkR :: a -> Set a -> Set a -> Set a +linkR x l r = case l of + Tip -> insertMin x r + Bin lsz _ _ _ -> linkR_ x lsz l r + +linkR_ :: a -> Int -> Set a -> Set a -> Set a +linkR_ x !lsz l r = case r of + Bin rsz rx rl rr + | delta*lsz < rsz -> balanceL rx (linkR_ x lsz l rl) rr + | otherwise -> Bin (1+lsz+rsz) x l r + Tip -> Bin (1+lsz) x l Tip + -- insertMin and insertMax don't perform potentially expensive comparisons. -insertMax,insertMin :: a -> Set a -> Set a +-- @since 0.8.1 +insertMax :: a -> Set a -> Set a insertMax x t = case t of Tip -> singleton x Bin _ y l r -> balanceR y l (insertMax x r) +-- @since 0.8.1 +insertMin :: a -> Set a -> Set a insertMin x t = case t of Tip -> singleton x @@ -1737,20 +1819,36 @@ -> balanceL y (insertMin x l) r {-------------------------------------------------------------------- - [merge l r]: merges two trees. + [link2 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 +-- @since 0.8.1 +link2 :: Set a -> Set a -> Set a +link2 Tip r = r +link2 l Tip = l +link2 l@(Bin lsz lx ll lr) r@(Bin rsz rx rl rr) + | delta*lsz < rsz = balanceL rx (link2R_ lsz l rl) rr + | delta*rsz < lsz = balanceR lx ll (link2L_ lr rsz r) + | otherwise = glue l r +link2L_ :: Set a -> Int -> Set a -> Set a +link2L_ l !rsz r = case l of + Bin lsz lx ll lr + | delta*rsz < lsz -> balanceR lx ll (link2L_ lr rsz r) + | otherwise -> glue l r + Tip -> r + +link2R_ :: Int -> Set a -> Set a -> Set a +link2R_ !lsz l r = case r of + Bin rsz rx rl rr + | delta*lsz < rsz -> balanceL rx (link2R_ lsz l rl) rr + | otherwise -> glue l r + Tip -> l + {-------------------------------------------------------------------- [glue l r]: glues two trees together. Assumes that [l] and [r] are already balanced with respect to each other. --------------------------------------------------------------------} +-- @since 0.8.1 glue :: Set a -> Set a -> Set a glue Tip r = r glue l Tip = l @@ -1760,8 +1858,9 @@ -- | \(O(\log n)\). Delete and find the minimal element. -- --- > deleteFindMin set = (findMin set, deleteMin set) - +-- Calls 'error' if the set is empty. +-- +-- __Note__: This function is partial. Prefer 'minView'. deleteFindMin :: Set a -> (a,Set a) deleteFindMin t | Just r <- minView t = r @@ -1769,7 +1868,9 @@ -- | \(O(\log n)\). Delete and find the maximal element. -- --- > deleteFindMax set = (findMax set, deleteMax set) +-- Calls 'error' if the set is empty. +-- +-- __Note__: This function is partial. Prefer 'maxView'. deleteFindMax :: Set a -> (a,Set a) deleteFindMax t | Just r <- maxView t = r @@ -1873,7 +1974,7 @@ -- 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. +-- Functions balanceL and balanceR are specialized versions of balance. -- balanceL only checks whether the left subtree is too big, -- balanceR only checks whether the right subtree is too big. @@ -1891,6 +1992,7 @@ -- balanceL is called when left subtree might have been inserted to or when -- right subtree might have been deleted from. +-- @since 0.8.1 balanceL :: a -> Set a -> Set a -> Set a balanceL x l r = case (l, r) of (Bin ls _ _ _, Bin rs _ _ _) @@ -1921,6 +2023,7 @@ -- balanceR is called when right subtree might have been inserted to or when -- left subtree might have been deleted from. +-- @since 0.8.1 balanceR :: a -> Set a -> Set a -> Set a balanceR x l r = case (l, r) of (Bin ls _ _ _, Bin rs _ _ _) @@ -2083,12 +2186,13 @@ newtype MergeSet a = MergeSet { getMergeSet :: Set a } instance Semigroup (MergeSet a) where - MergeSet xs <> MergeSet ys = MergeSet (merge xs ys) + MergeSet xs <> MergeSet ys = MergeSet (link2 xs ys) instance Monoid (MergeSet a) where mempty = MergeSet empty - +#if !MIN_VERSION_base(4,11,0) mappend = (<>) +#endif -- | \(O(n+m)\). Calculate the disjoint union of two sets. -- @@ -2103,9 +2207,345 @@ -- -- @since 0.5.11 disjointUnion :: Set a -> Set b -> Set (Either a b) -disjointUnion as bs = merge (mapMonotonic Left as) (mapMonotonic Right bs) +disjointUnion as bs = link2 (mapMonotonic Left as) (mapMonotonic Right bs) {-------------------------------------------------------------------- + Merging Sets +--------------------------------------------------------------------} + +-- | A tactic for dealing with elements present in one set but not the other in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @WhenMissing f a@ is an abstract representation of a +-- function of type @a -> f Bool@. +-- +-- @since 0.8.1 +data WhenMissing f a = WhenMissing + { missingSubtree :: Set a -> f (Set a) + , missingElem :: a -> f Bool + } + +-- | A tactic for dealing with elements present in one set but not the other in +-- 'merge'. +-- +-- A tactic of type @SimpleWhenMissing a@ is an abstract representation +-- of a function of type @a -> Bool@. +-- +-- @since 0.8.1 +type SimpleWhenMissing = WhenMissing Identity + +-- | Along with 'filterAMissing', witnesses the isomorphism between +-- @WhenMissing f a@ and @a -> f Bool@. +-- +-- @since 0.8.1 +runWhenMissing :: WhenMissing f a -> a -> f Bool +runWhenMissing = missingElem + +-- | A tactic for dealing with elements present in both sets in 'merge' or +-- 'mergeA'. +-- +-- A tactic of type @WhenMatched f a@ is an abstract representation of a +-- function of type @a -> f Bool@. +-- +-- @since 0.8.1 +newtype WhenMatched f a = WhenMatched { matchedElem :: a -> f Bool } + +-- | A tactic for dealing with elements present in both sets in 'merge'. +-- +-- A tactic of type @SimpleWhenMatched a@ is an abstract representation of a +-- function of type @a -> Bool@. +-- +-- @since 0.8.1 +type SimpleWhenMatched = WhenMatched Identity + +-- | Along with 'filterAMatched', witnesses the isomorphism between +-- @WhenMatched f a@ and @a -> f Bool@. +-- +-- @since 0.8.1 +runWhenMatched :: WhenMatched f a -> a -> f Bool +runWhenMatched = matchedElem + +-- | When an element is found in both sets, drop the element. +-- +-- @since 0.8.1 +dropMatched :: Applicative f => WhenMatched f a +dropMatched = WhenMatched (\_ -> pure False) +{-# INLINE dropMatched #-} + +-- | When an element is found in both sets, keep the element. +-- +-- @since 0.8.1 +preserveMatched :: Applicative f => WhenMatched f a +preserveMatched = WhenMatched (\_ -> pure True) +{-# INLINE preserveMatched #-} + +-- | When an element is found in both sets, choose whether to keep the element +-- in the merged set. +-- +-- @since 0.8.1 +filterMatched :: Applicative f => (a -> Bool) -> WhenMatched f a +filterMatched f = WhenMatched (pure . f) +{-# INLINE filterMatched #-} + +-- | When an element is found in both sets, choose whether to keep the element +-- in the merged set. +-- +-- @since 0.8.1 +filterAMatched :: (a -> f Bool) -> WhenMatched f a +filterAMatched = WhenMatched + +-- | Create a @WhenMissing@ from two functions. +-- +-- @whenMissing@ must be called with two functions @f@ and @g@ such that +-- @g = 'filterA' f@. @g@ may be a more efficient way of applying @f@ to all +-- elements in a @Set@. +-- +-- __Warning__: It is the caller's responsibility to ensure the above property. +-- +-- === __Examples__ +-- +-- @ +-- preserveMissing :: Applicative f => WhenMissing f a +-- preserveMissing = whenMissing f g +-- where +-- f _x = pure True +-- g s = pure s +-- -- Note that this satisfies g = filterA f +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- For a usage of this, see examples on mergeA +-- isEmpty :: WhenMissing (Const All) a +-- isEmpty = whenMissing f g +-- where +-- f _x = Const (All False) +-- g s = Const (All (null s)) +-- -- Note that this satisfies g = filterA f +-- @ +-- +-- @since 0.8.1 +whenMissing :: (a -> f Bool) -> (Set a -> f (Set a)) -> WhenMissing f a +whenMissing = flip WhenMissing + +-- | Drop all the elements that are missing from the other set. +-- +-- @ +-- dropMissing :: 'SimpleWhenMissing' a +-- @ +-- +-- > dropMissing = filterMissing (\_ -> False) +-- +-- but @dropMissing@ is more efficient. +-- +-- @since 0.8.1 +dropMissing :: Applicative f => WhenMissing f a +dropMissing = WhenMissing + { missingSubtree = \_ -> pure Tip + , missingElem = \_ -> pure False + } +{-# INLINE dropMissing #-} + +-- | Preserve the elements that are missing from the other set. +-- +-- @ +-- preserveMissing :: 'SimpleWhenMissing' a +-- @ +-- +-- > preserveMissing = filterMissing (\_ -> True) +-- +-- but @preserveMissing@ is more efficient. +-- +-- @since 0.8.1 +preserveMissing :: Applicative f => WhenMissing f a +preserveMissing = WhenMissing + { missingSubtree = pure + , missingElem = \_ -> pure True + } +{-# INLINE preserveMissing #-} + +-- | Filter the elements that are missing from the other set. +-- +-- @ +-- filterMissing :: (a -> Bool) -> 'SimpleWhenMissing' a +-- @ +-- +-- @since 0.8.1 +filterMissing :: Applicative f => (a -> Bool) -> WhenMissing f a +filterMissing f = WhenMissing + { missingSubtree = pure . filter f + , missingElem = pure . f + } +{-# INLINE filterMissing #-} + +-- | Filter the elements that are missing from the other set using some +-- 'Applicative' action. +-- +-- @since 0.8.1 +filterAMissing :: Applicative f => (a -> f Bool) -> WhenMissing f a +filterAMissing f = WhenMissing + { missingSubtree = filterA f + , missingElem = f + } +{-# INLINE filterAMissing #-} + +-- | Merge two sets. +-- +-- 'merge' takes two 'SimpleWhenMissing' tactics, a 'SimpleWhenMatched' tactic, +-- and two sets. It uses the tactics to merge the sets. +-- +-- Consider +-- +-- @ +-- merge (filterMissing g1) (filterMissing g2) (filterMatched f) s1 s2 +-- @ +-- +-- @ +-- g1 = (==2) +-- g2 = (==3) +-- f = (==6) +-- s1 = [2, 4, 6, 8, 10, 12] +-- s2 = [3, 6, 9, 12] +-- @ +-- +-- 'merge' will pass the elements to @g1@, @g2@, or @f@ as appropriate, +-- producing a @Bool@ for each element. +-- +-- @ +-- m1: [ 2, 4, 6, 8, 10 12] +-- m2: [ 3, 6, 9, 12] +-- result: [g1 2, g2 3, g1 4, f 6, g1 8, g2 9, g1 10, f 12] +-- = [True, True, False, True, False, False, False, False] +-- @ +-- +-- The result set contains the element for which we have @True@. +-- +-- >>> merge (filterMissing g1) (filterMissing g2) (filterMatched f) s1 s2 +-- fromList [2,3,6] +-- +-- The other tactics below are optimizations or simplifications of +-- 'filterMissing' for special cases. Most importantly, +-- +-- * 'dropMissing' drops all elements. +-- * 'preserveMissing' leaves all elements alone. +-- +-- When 'merge' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should typically use 'merge' +-- to define your custom combining functions. +-- +-- @since 0.8.1 +merge + :: Ord a + => SimpleWhenMissing a -- ^ What to do with elements in @s1@ but not @s2@ + -> SimpleWhenMissing a -- ^ What to do with elements in @s2@ but not @s1@ + -> SimpleWhenMatched a -- ^ What to do with elements in both @s1@ and @s2@ + -> Set a -- ^ Set @s1@ + -> Set a -- ^ Set @s2@ + -> Set a +merge g1 g2 f = \s1 s2 -> runIdentity (mergeA g1 g2 f s1 s2) +{-# INLINE merge #-} + +-- | An applicative version of 'merge'. +-- +-- 'mergeA' takes two 'WhenMissing' tactics, a 'WhenMatched' tactic, and two +-- sets. It uses the tactics to merge the sets. +-- +-- Behaves just like 'merge' while allowing @Applicative@ effects. Effects are +-- performed in increasing order of elements. +-- +-- Consider, +-- +-- @ +-- mergeA (filterAMissing g1) (filterAMissing g2) (filterAMatched f) s1 s2 +-- @ +-- +-- @ +-- g1 x = (x == 2) <$ putStrLn ("g1 " ++ show x) +-- g2 x = (x == 3) <$ putStrLn ("g2 " ++ show x) +-- f x = (x == 6) <$ putStrLn ("f " ++ show x) +-- s1 = [2, 4, 6, 8, 10, 12] +-- s2 = [3, 6, 9, 12] +-- @ +-- +-- As with 'merge', the result set is @[2,3,6]@. Additionally, @g1@, @g2@, and +-- @f@ perform @IO@ effects, printing the elements in increasing order. +-- +-- >>> mergeA (filterAMissing g1) (filterAMissing g2) (filterAMatched f) s1 s2 +-- g1 2 +-- g2 3 +-- g1 4 +-- f 6 +-- g1 8 +-- g2 9 +-- g1 10 +-- f 12 +-- fromList [2,3,6] +-- +-- When 'mergeA' is given three arguments, it is inlined at the call +-- site. To prevent excessive inlining, you should generally only use +-- 'mergeA' to define custom combining functions. +-- +-- === __Examples__ +-- +-- @ +-- data Pair a = Pair !a !a deriving Functor +-- +-- instance Applicative Pair where +-- pure x = Pair x x +-- liftA2 f (Pair x1 y1) (Pair x2 y2) = Pair (f x1 x2) (f y1 y2) +-- +-- -- | Calculate the union and intersection of two sets. +-- unionIntersection :: Ord a => Set a -> Set a -> (Set a, Set a) +-- unionIntersection m1 m2 = +-- case mergeA preserveAndDropMissing preserveAndDropMissing 'preserveMatched' m1 m2 of +-- Pair mu mi -> (mu, mi) +-- where +-- -- use Pair to build the union and intersection together +-- preserveAndDropMissing = 'whenMissing' (\\_x -> Pair True False) (\\s -> Pair s empty) +-- @ +-- +-- @ +-- import Data.Functor.Const (Const(..)) +-- import Data.Monoid (All(..)) +-- +-- -- | Whether the first set is a subset of the second set. +-- isSubsetOf :: Ord a => Set a -> Set a -> Bool +-- isSubsetOf m1 m2 = +-- getAll (getConst (mergeA isEmpty 'dropMissing' 'dropMatched' m1 m2)) +-- where +-- isEmpty = 'whenMissing' (\\_x -> Const (All False)) (\\s -> Const (All (null s))) +-- @ +-- +-- @since 0.8.1 +mergeA + :: (Applicative f, Ord a) + => WhenMissing f a -- ^ What to do with elements in @s1@ but not @s2@ + -> WhenMissing f a -- ^ What to do with elements in @s2@ but not @s1@ + -> WhenMatched f a -- ^ What to do with elements in both @s1@ and @s2@ + -> Set a -- ^ Set @s1@ + -> Set a -- ^ Set @s2@ + -> f (Set a) +mergeA + WhenMissing{missingSubtree = g1t, missingElem = g1k} + WhenMissing{missingSubtree = g2t} + WhenMatched{matchedElem = f} = go + where + go t1 Tip = g1t t1 + go Tip t2 = g2t t2 + go (Bin _ x1 l1 r1) t2 = case splitMember x1 t2 of + (l2, found, r2) + | found -> liftA3 doLink l1l2 (f x1) r1r2 + | otherwise -> liftA3 doLink l1l2 (g1k x1) r1r2 + where + doLink l' True r' = link x1 l' r' + doLink l' False r' = link2 l' r' + l1l2 = go l1 l2 + r1r2 = go r1 r2 +{-# INLINE mergeA #-} + +{-------------------------------------------------------------------- Debugging --------------------------------------------------------------------} -- | \(O(n \log n)\). Show the tree that implements the set. The tree is shown @@ -2266,8 +2706,8 @@ -- done in O(1) using `Bin`. The final linking of the stack is done in O(log n) -- using `link` (proof below). The total time is thus O(n). -- --- Additionally, the implemention is written using foldl' over the input list, --- which makes it participate as a good consumer in list fusion. +-- Additionally, the implementation is written using foldl' over the input +-- list, which makes it participate as a good consumer in list fusion. -- -- fromDistinctDescList is implemented similarly, adapted for left and right -- sides being swapped. @@ -2286,3 +2726,34 @@ -- = O(\sum_{i=2}^m k_i - k_{i-1}) -- = O(k_m - k_1) -- = O(log n) + +-- Note [alterF implementation] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- +-- When implementing alterF, there are various costs to consider: +-- +-- 1. The number of times we travel down the tree to the target +-- 2. The number of times we travel back up the tree +-- 3. The number of element comparisons we perform +-- 4. The number of times we apply fmap +-- 5. The amount of allocations we perform +-- +-- The weights of these further depend on the Ord instance of the element type, +-- the Functor used, and GHC optimizations. Unfortunately, there is no clear +-- winning implementation which performs best in all situations. The current +-- implementation is chosen for its good characteristics: +-- +-- 1. It travels down the tree once if the membership is demanded, and once more +-- if the changed Set is demanded. +-- 2. It travels up the tree once if the changed Set is demanded. +-- 3. It compares the given element against stored elements on the path to the +-- target at most once. +-- 4. It applies fmap exactly once. +-- 5. It allocates the changed Set if it is demanded. It does not allocate any +-- intermediate structure. +-- +-- Note that 2,3,4,5 are all optimal. As for 1, implementations that travel down +-- the tree at most once are possible, but they are worse in at least one of the +-- other categories. +-- See https://github.com/haskell/containers/issues/1209#issuecomment-4530740850 +-- for some other implementations and benchmarks. diff --git a/src/Data/Set/Merge.hs b/src/Data/Set/Merge.hs new file mode 100644 --- /dev/null +++ b/src/Data/Set/Merge.hs @@ -0,0 +1,53 @@ +{-# LANGUAGE CPP #-} +#ifdef __GLASGOW_HASKELL__ +{-# LANGUAGE Safe #-} +#endif + +-- | This module defines an API for writing functions that merge two sets. The key +-- functions are 'merge' and 'mergeA'. Each of these can be used with several +-- different \"merge tactics\". +-- +-- @since 0.8.1 +module Data.Set.Merge + ( + -- ** Simple merge tactic types + SimpleWhenMissing + , SimpleWhenMatched + + -- ** General combining function + , merge + + -- *** @WhenMissing@ tactics + , dropMissing + , preserveMissing + , filterMissing + + -- *** @WhenMatched@ tactics + , dropMatched + , preserveMatched + , filterMatched + + -- ** Applicative merge tactic types + , WhenMissing + , WhenMatched + + -- ** Applicative general combining function + , mergeA + + -- *** @WhenMissing@ tactics + -- | The tactics described for 'merge' work for 'mergeA' as well. + -- Furthermore, the following are available. + , filterAMissing + , whenMissing + + -- *** @WhenMatched@ tactics + -- | The tactics described for 'merge' work for 'mergeA' as well. + -- Furthermore, the following are available. + , filterAMatched + + -- ** Miscellaneous tactic functions + , runWhenMissing + , runWhenMatched + ) where + +import Data.Set.Internal diff --git a/src/Data/Tree.hs b/src/Data/Tree.hs --- a/src/Data/Tree.hs +++ b/src/Data/Tree.hs @@ -1,5 +1,4 @@ {-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} {-# LANGUAGE CPP #-} #if __GLASGOW_HASKELL__ {-# LANGUAGE DeriveDataTypeable #-} @@ -8,8 +7,6 @@ {-# LANGUAGE Trustworthy #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.Tree @@ -75,9 +72,13 @@ import Data.Data (Data) import GHC.Generics (Generic, Generic1) import qualified GHC.Exts +# if __GLASGOW_HASKELL__ >= 914 +import Language.Haskell.TH.Lift (Lift) +# else import Language.Haskell.TH.Syntax (Lift) -- See Note [ Template Haskell Dependencies ] import Language.Haskell.TH () +# endif #endif import Control.Monad.Zip (MonadZip (..)) @@ -210,10 +211,7 @@ foldr f z = \t -> go t z -- Use a lambda to allow inlining with two arguments where - go (Node x ts) = f x . foldr (\t k -> go t . k) id ts - -- This is equivalent to the following simpler definition, but has been found to optimize - -- better in benchmarks: - -- go (Node x ts) z' = f x (foldr go z' ts) + go (Node x ts) z' = f x (foldrTreeList go z' ts) {-# INLINE foldr #-} foldl' f = go @@ -242,6 +240,17 @@ product = foldlMap1' id (*) {-# INLINABLE product #-} +-- This is the same as List's foldr, but unlike GHC's implementation the z is +-- passed along in go instead of go closing over it. +-- When folding over a Tree this avoids a closure per Node, which results in +-- significant reductions in time and allocations according to benchmarks. +foldrTreeList :: (Tree a -> b -> b) -> b -> [Tree a] -> b +foldrTreeList f = go + where + go z [] = z + go z (t:ts) = f t (go z ts) +{-# INLINE foldrTreeList #-} + #if MIN_VERSION_base(4,18,0) -- | Folds in pre-order. -- @@ -497,10 +506,12 @@ (a, bs) <- f b ts <- unfoldForestM f bs return (Node a ts) +{-# INLINABLE unfoldTreeM #-} -- | Monadic forest builder, in depth-first order. unfoldForestM :: Monad m => (b -> m (a, [b])) -> [b] -> m ([Tree a]) unfoldForestM f = Prelude.mapM (unfoldTreeM f) +{-# INLINABLE unfoldForestM #-} -- | Monadic tree builder, in breadth-first order. -- @@ -515,6 +526,7 @@ getElement xs = case viewl xs of x :< _ -> x EmptyL -> error "unfoldTreeM_BF" +{-# INLINABLE unfoldTreeM_BF #-} -- | Monadic forest builder, in breadth-first order. -- @@ -525,6 +537,7 @@ -- by Chris Okasaki, /ICFP'00/. unfoldForestM_BF :: Monad m => (b -> m (a, [b])) -> [b] -> m ([Tree a]) unfoldForestM_BF f = liftM toList . unfoldForestQ f . fromList +{-# INLINABLE unfoldForestM_BF #-} -- Takes a sequence (queue) of seeds and produces a sequence (reversed queue) of -- trees of the same length. @@ -542,6 +555,7 @@ splitOnto as (_:bs) q = case viewr q of q' :> a -> splitOnto (a:as) bs q' EmptyR -> error "unfoldForestQ" +{-# INLINABLE unfoldForestQ #-} -- | \(O(n)\). The leaves of the tree in left-to-right order. -- @@ -570,7 +584,7 @@ #ifdef __GLASGOW_HASKELL__ leaves t = GHC.Exts.build $ \cons nil -> let go (Node x []) z = cons x z - go (Node _ ts) z = foldr go z ts + go (Node _ ts) z = foldrTreeList go z ts in go t nil {-# INLINE leaves #-} -- Inline for list fusion #else @@ -606,8 +620,9 @@ edges :: Tree a -> [(a, a)] #ifdef __GLASGOW_HASKELL__ edges (Node x0 ts0) = GHC.Exts.build $ \cons nil -> - let go p = foldr (\(Node x ts) z -> cons (p, x) (go x z ts)) - in go x0 nil ts0 + let go _ [] z = z + go p (Node x ts : ts') z = cons (p, x) (go x ts (go p ts' z)) + in go x0 ts0 nil {-# INLINE edges #-} -- Inline for list fusion #else edges (Node x0 ts0) = @@ -691,6 +706,25 @@ -- | A newtype over 'Tree' that folds and traverses in post-order. -- +-- ==== __@Foldable@ examples__ +-- +-- >>> import Data.Foldable (toList) +-- >>> toList $ PostOrder $ Node 1 [Node 2 [Node 3 [], Node 4 []], Node 5 []] +-- [3,4,2,5,1] +-- +-- @foldr@ produces elements incrementally, inspecting the structure of the +-- @Tree@ just as much as necessary. +-- +-- >>> take 3 $ foldr (:) [] $ PostOrder $ Node 1 ([Node 2 [Node 3 [], Node 4 []]] ++ undefined) +-- [3,4,2] +-- +-- @foldl@ also produces elements incrementally. +-- +-- >>> foldl (flip (:)) [] $ PostOrder $ Node 1 [Node 2 [Node 3 [], Node 4 []], Node 5 []] +-- [1,5,2,4,3] +-- >>> take 4 $ foldl (flip (:)) [] $ PostOrder $ Node 1 [Node 2 [undefined, Node 4 []], Node 5 []] +-- [1,5,2,4] +-- -- @since 0.8 newtype PostOrder a = PostOrder { unPostOrder :: Tree a } #ifdef __GLASGOW_HASKELL__ @@ -722,7 +756,7 @@ foldr f z0 = \(PostOrder t) -> go t z0 -- Use a lambda to inline with two arguments where - go (Node x ts) z = foldr go (f x z) ts + go (Node x ts) z = foldrTreeList go (f x z) ts {-# INLINE foldr #-} foldl' f z0 = \(PostOrder t) -> go z0 t -- Use a lambda to inline with two arguments @@ -732,9 +766,25 @@ in f z' x {-# INLINE foldl' #-} + foldl f z0 = -- Inline with two arguments + \(PostOrder t) -> go z0 t + where + go z (Node x ts) = f (Foldable.foldl go z ts) x + {-# INLINE foldl #-} + + foldr' f z0 = -- Inline with two arguments + \(PostOrder t) -> go t z0 + where + go (Node x ts) !z = + let !z' = f x z + in foldrTreeList go z' ts + {-# INLINE foldr' #-} + foldr1 = foldrMap1PostOrder id + {-# INLINE foldr1 #-} foldl1 = foldlMap1PostOrder id + {-# INLINE foldl1 #-} null _ = False {-# INLINE null #-} @@ -777,7 +827,7 @@ where go (Node x []) z = x :| z go (Node x (t:ts)) z = - go t (foldr (\t' z' -> foldr (:) z' (PostOrder t')) (x:z) ts) + go t (foldrTreeList (\t' z' -> foldr (:) z' (PostOrder t')) (x:z) ts) maximum = Foldable.maximum {-# INLINABLE maximum #-} @@ -786,10 +836,18 @@ {-# INLINABLE minimum #-} foldrMap1 = foldrMap1PostOrder + {-# INLINE foldrMap1 #-} foldlMap1' = foldlMap1'PostOrder + {-# INLINE foldlMap1' #-} foldlMap1 = foldlMap1PostOrder + {-# INLINE foldlMap1 #-} + + foldrMap1' f g = -- Inline with two arguments + \(PostOrder (Node x ts)) -> + foldr (\t !z -> Foldable.foldr' g z (PostOrder t)) (f x) ts + {-# INLINE foldrMap1' #-} #endif foldrMap1PostOrder :: (a -> b) -> (a -> b -> b) -> PostOrder a -> b diff --git a/src/Utils/Containers/Internal/BitQueue.hs b/src/Utils/Containers/Internal/BitQueue.hs --- a/src/Utils/Containers/Internal/BitQueue.hs +++ b/src/Utils/Containers/Internal/BitQueue.hs @@ -1,7 +1,4 @@ -{-# LANGUAGE CPP #-} {-# LANGUAGE BangPatterns #-} - -#include "containers.h" ----------------------------------------------------------------------------- -- | diff --git a/src/Utils/Containers/Internal/BitUtil.hs b/src/Utils/Containers/Internal/BitUtil.hs --- a/src/Utils/Containers/Internal/BitUtil.hs +++ b/src/Utils/Containers/Internal/BitUtil.hs @@ -1,10 +1,7 @@ {-# LANGUAGE CPP #-} #ifdef __GLASGOW_HASKELL__ {-# LANGUAGE MagicHash #-} -{-# LANGUAGE Trustworthy #-} #endif - -#include "containers.h" ----------------------------------------------------------------------------- -- | diff --git a/src/Utils/Containers/Internal/EqOrdUtil.hs b/src/Utils/Containers/Internal/EqOrdUtil.hs --- a/src/Utils/Containers/Internal/EqOrdUtil.hs +++ b/src/Utils/Containers/Internal/EqOrdUtil.hs @@ -7,7 +7,7 @@ #if !MIN_VERSION_base(4,11,0) import Data.Semigroup (Semigroup(..)) #endif -import Utils.Containers.Internal.StrictPair +import Utils.Containers.Internal.Strict (StrictPair(..)) newtype EqM a = EqM { runEqM :: a -> StrictPair Bool a } diff --git a/src/Utils/Containers/Internal/State.hs b/src/Utils/Containers/Internal/State.hs --- a/src/Utils/Containers/Internal/State.hs +++ b/src/Utils/Containers/Internal/State.hs @@ -1,5 +1,3 @@ -{-# LANGUAGE CPP #-} -#include "containers.h" {-# OPTIONS_HADDOCK hide #-} -- | A clone of Control.Monad.State.Strict. diff --git a/src/Utils/Containers/Internal/Strict.hs b/src/Utils/Containers/Internal/Strict.hs new file mode 100644 --- /dev/null +++ b/src/Utils/Containers/Internal/Strict.hs @@ -0,0 +1,22 @@ +-- | Simple strict types for internal use. +module Utils.Containers.Internal.Strict + ( StrictPair(..) + , toPair + , StrictTriple(..) + ) where + +-- | The same as a regular Haskell pair, but +-- +-- @ +-- (x :*: _|_) = (_|_ :*: y) = _|_ +-- @ +data StrictPair a b = !a :*: !b + +infixr 1 :*: + +-- | Convert a strict pair to a standard pair. +toPair :: StrictPair a b -> (a, b) +toPair (x :*: y) = (x, y) +{-# INLINE toPair #-} + +data StrictTriple a b c = TripleS !a !b !c diff --git a/src/Utils/Containers/Internal/StrictMaybe.hs b/src/Utils/Containers/Internal/StrictMaybe.hs deleted file mode 100644 --- a/src/Utils/Containers/Internal/StrictMaybe.hs +++ /dev/null @@ -1,29 +0,0 @@ -{-# LANGUAGE CPP #-} - -#include "containers.h" - -{-# OPTIONS_HADDOCK hide #-} --- | Strict 'Maybe' - -module Utils.Containers.Internal.StrictMaybe (MaybeS (..), maybeS, toMaybe, toMaybeS) where -#ifdef __MHS__ -import Data.Foldable -#endif - -data MaybeS a = NothingS | JustS !a - -instance Foldable MaybeS where - foldMap _ NothingS = mempty - foldMap f (JustS a) = f a - -maybeS :: r -> (a -> r) -> MaybeS a -> r -maybeS n _ NothingS = n -maybeS _ j (JustS a) = j a - -toMaybe :: MaybeS a -> Maybe a -toMaybe NothingS = Nothing -toMaybe (JustS a) = Just a - -toMaybeS :: Maybe a -> MaybeS a -toMaybeS Nothing = NothingS -toMaybeS (Just a) = JustS a diff --git a/src/Utils/Containers/Internal/StrictPair.hs b/src/Utils/Containers/Internal/StrictPair.hs deleted file mode 100644 --- a/src/Utils/Containers/Internal/StrictPair.hs +++ /dev/null @@ -1,24 +0,0 @@ -{-# LANGUAGE CPP #-} -#ifdef __GLASGOW_HASKELL__ -{-# LANGUAGE Safe #-} -#endif - -#include "containers.h" - --- | A strict pair - -module Utils.Containers.Internal.StrictPair (StrictPair(..), toPair) where - --- | The same as a regular Haskell pair, but --- --- @ --- (x :*: _|_) = (_|_ :*: y) = _|_ --- @ -data StrictPair a b = !a :*: !b - -infixr 1 :*: - --- | Convert a strict pair to a standard pair. -toPair :: StrictPair a b -> (a, b) -toPair (x :*: y) = (x, y) -{-# INLINE toPair #-}