diff --git a/containers-tests/benchmarks/IntMap.hs b/containers-tests/benchmarks/IntMap.hs index 8bfc24271..c3d2210ff 100644 --- a/containers-tests/benchmarks/IntMap.hs +++ b/containers-tests/benchmarks/IntMap.hs @@ -13,9 +13,12 @@ import Prelude hiding (lookup) main = do let m = M.fromAscList elems :: M.IntMap Int evaluate $ rnf [m] + evaluate $ rnf missKeys defaultMain - [ bench "lookup" $ whnf (lookup keys) m - , bench "insert" $ whnf (ins elems) M.empty + [ bench "size" $ whnf M.size m + , bench "lookup hit" $ whnf (lookup keys) m + , bench "lookup miss" $ whnf (lookup missKeys) m + , bench "insert empty" $ whnf (ins elems) M.empty , bench "insertWith empty" $ whnf (insWith elems) M.empty , bench "insertWith update" $ whnf (insWith elems) m , bench "insertWith' empty" $ whnf (insWith' elems) M.empty @@ -28,13 +31,18 @@ main = do , bench "insertLookupWithKey update" $ whnf (insLookupWithKey elems) m , bench "map" $ whnf (M.map (+ 1)) m , bench "mapWithKey" $ whnf (M.mapWithKey (+)) m - , bench "foldlWithKey" $ whnf (ins elems) m + , bench "foldlWithKey" $ whnf (M.foldlWithKey consPairL []) m , bench "foldlWithKey'" $ whnf (M.foldlWithKey' sum 0) m - , bench "foldrWithKey" $ whnf (M.foldrWithKey consPair []) m - , bench "delete" $ whnf (del keys) m - , bench "update" $ whnf (upd keys) m - , bench "updateLookupWithKey" $ whnf (upd' keys) m - , bench "alter" $ whnf (alt keys) m + , bench "foldrWithKey" $ whnf (M.foldrWithKey consPairR []) m + , bench "foldrWithKey'" $ whnf (M.foldrWithKey' sum 0) m + , bench "delete hit" $ whnf (del keys) m + , bench "delete miss" $ whnf (del missKeys) m + , bench "update hit" $ whnf (upd keys) m + , bench "update miss" $ whnf (upd missKeys) m + , bench "updateLookupWithKey hit" $ whnf (upd' keys) m + , bench "updateLookupWithKey miss" $ whnf (upd' missKeys) m + , bench "alter hit" $ whnf (alt keys) m + , bench "alter miss" $ whnf (alt missKeys) m , bench "mapMaybe" $ whnf (M.mapMaybe maybeDel) m , bench "mapMaybeWithKey" $ whnf (M.mapMaybeWithKey (const maybeDel)) m , bench "fromList" $ whnf M.fromList elems @@ -45,10 +53,12 @@ main = do ] where elems = zip keys values - keys = [1..2^12] - values = [1..2^12] + keys = [1,3..2^13] + missKeys = [0,2..2^13] + values = [1,3..2^13] sum k v1 v2 = k + v1 + v2 - consPair k v xs = (k, v) : xs + consPairL xs k v = (k, v) : xs + consPairR k v xs = (k, v) : xs add3 :: Int -> Int -> Int -> Int add3 x y z = x + y + z diff --git a/containers-tests/benchmarks/LookupGE/IntMap.hs b/containers-tests/benchmarks/LookupGE/IntMap.hs index e18d669ca..7c53a2cfd 100644 --- a/containers-tests/benchmarks/LookupGE/IntMap.hs +++ b/containers-tests/benchmarks/LookupGE/IntMap.hs @@ -34,9 +34,7 @@ main = do , \(n,fun) -> bench (n++" !far") $ nf (fge2 fun odds) m_large ] funs1 = [ ("GE split", M.lookupGE1) - , ("GE Craig", M.lookupGE2) - , ("GE Twan", M.lookupGE3) - , ("GE Milan", M.lookupGE4) ] + , ("GE direct", M.lookupGE2) ] fge :: (Int -> M.IntMap Int -> Maybe (Int,Int)) -> [Int] -> M.IntMap Int -> (Int,Int) fge fun xs m = foldl' (\n k -> fromMaybe n (fun k m)) (0,0) xs diff --git a/containers-tests/benchmarks/LookupGE/LookupGE_IntMap.hs b/containers-tests/benchmarks/LookupGE/LookupGE_IntMap.hs index ff849b1d2..c9512ce64 100644 --- a/containers-tests/benchmarks/LookupGE/LookupGE_IntMap.hs +++ b/containers-tests/benchmarks/LookupGE/LookupGE_IntMap.hs @@ -10,63 +10,8 @@ lookupGE1 k m = (_,Just v,_) -> Just (k,v) (_,Nothing,r) -> findMinMaybe r -lookupGE2 :: Key -> IntMap a -> Maybe (Key,a) -lookupGE2 k t = case t of - Bin _ m l r | m < 0 -> if k >= 0 - then go l - else case go r of - Nothing -> Just $ findMin l - justx -> justx - _ -> go t - where - go (Bin p m l r) - | nomatch k p m = if k < p - then Just $ findMin l - else Nothing - | zero k m = case go l of - Nothing -> Just $ findMin r - justx -> justx - | otherwise = go r - go (Tip ky y) - | k > ky = Nothing - | otherwise = Just (ky, y) - go Nil = Nothing -lookupGE3 :: Key -> IntMap a -> Maybe (Key,a) -lookupGE3 k t = k `seq` case t of - Bin _ m l r | m < 0 -> if k >= 0 - then go Nothing l - else go (Just (findMin l)) r - _ -> go Nothing t - where - go def (Bin p m l r) - | nomatch k p m = if k < p then Just $ findMin l else def - | zero k m = go (Just $ findMin r) l - | otherwise = go def r - go def (Tip ky y) - | k > ky = def - | otherwise = Just (ky, y) - go def Nil = def - -lookupGE4 :: Key -> IntMap a -> Maybe (Key,a) -lookupGE4 k t = k `seq` case t of - Bin _ m l r | m < 0 -> if k >= 0 then go Nil l - else go l r - _ -> go Nil t - where - go def (Bin p m l r) - | nomatch k p m = if k < p then fMin l else fMin def - | zero k m = go r l - | otherwise = go def r - go def (Tip ky y) - | k > ky = fMin def - | otherwise = Just (ky, y) - go def Nil = fMin def - - fMin :: IntMap a -> Maybe (Key, a) - fMin Nil = Nothing - fMin (Tip ky y) = Just (ky, y) - fMin (Bin _ _ l _) = fMin l +lookupGE2 = lookupGE ------------------------------------------------------------------------------- -- Utilities @@ -85,10 +30,5 @@ findMinMaybe m prop_lookupGE12 :: Int -> [Int] -> Bool prop_lookupGE12 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE2 x m - -prop_lookupGE13 :: Int -> [Int] -> Bool -prop_lookupGE13 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE3 x m - -prop_lookupGE14 :: Int -> [Int] -> Bool -prop_lookupGE14 x xs = case fromList $ zip xs xs of m -> lookupGE1 x m == lookupGE4 x m #endif + diff --git a/containers-tests/benchmarks/SetOperations/SetOperations-IntMap.hs b/containers-tests/benchmarks/SetOperations/SetOperations-IntMap.hs index 036c82caa..4f5a1bd65 100644 --- a/containers-tests/benchmarks/SetOperations/SetOperations-IntMap.hs +++ b/containers-tests/benchmarks/SetOperations/SetOperations-IntMap.hs @@ -3,4 +3,12 @@ module Main where import Data.IntMap as C import SetOperations -main = benchmark (\xs -> fromList [(x, x) | x <- xs]) True [("union", C.union), ("difference", C.difference), ("intersection", C.intersection)] +main = benchmark (\xs -> fromList [(x, x) | x <- xs]) True + [ ("union", C.union) + , ("unionWith", C.unionWith (+)) + , ("difference", C.difference) + , ("differenceWith (keep)", C.differenceWith (\a b -> Just (a + b))) + , ("differenceWith (delete)", C.differenceWith (\_ _ -> Nothing)) + , ("intersection", C.intersection) + , ("intersectionWith", C.intersectionWith (+)) + ] diff --git a/containers-tests/benchmarks/bench-cmp.pl b/containers-tests/benchmarks/bench-cmp.pl index 52875ae87..10875c94f 100755 --- a/containers-tests/benchmarks/bench-cmp.pl +++ b/containers-tests/benchmarks/bench-cmp.pl @@ -10,6 +10,11 @@ my $l2 = <$f2>; $l1 eq $l2 or die "CSV files do not correspond -- $l1 and $l2"; +my $min = 1e50; +my $mult = 1.0; +my $max = 0.0; +my $count = 0; + while (defined($l1 = <$f1>)) { $l2 = <$f2>; @@ -17,8 +22,37 @@ my @parts2 = split /,/, $l2; $parts1[0] eq $parts2[0] or die "CSV files do not correspond -- $parts1[0] and $parts2[0]"; - printf "%s;%+7.2f%%;%.2e\n", $parts1[0], 100 * $parts2[1] / $parts1[1] - 100, $parts1[1]; + + my $factor; + if ($parts1[1] == 0) { + if ($parts2[1] == 0) { + $factor = 1; + } else { + $factor = 'inf'; + } + } else { + $factor = $parts2[1] / $parts1[1]; + } + $count = $count + 1; + $mult = $mult * $factor; + if ($factor > $max) { + $max = $factor; + } + if ($factor < $min) { + $min = $factor; + } + + if ($factor == 'inf') { + printf "%s;%.2e;%.2e\n", $parts1[0], $parts2[1], $parts1[1]; + } else { + printf "%s;%+7.2f%%;%.2e\n", $parts1[0], 100 * $factor - 100, $parts1[1]; + } } +printf ";\n"; +printf "Minimum;%+7.2f%%\n", ($min - 1.0)*100; +printf "Average;%+7.2f%%\n", (($mult ** (1.0 / $count)) - 1.0) * 100; +printf "Maximum;%+7.2f%%\n", ($max - 1.0) * 100; + close $f2; close $f1; diff --git a/containers-tests/containers-tests.cabal b/containers-tests/containers-tests.cabal index 5be41e070..c9cc16472 100644 --- a/containers-tests/containers-tests.cabal +++ b/containers-tests/containers-tests.cabal @@ -55,10 +55,10 @@ library Data.IntMap.Internal Data.IntMap.Internal.Debug Data.IntMap.Lazy + Data.IntMap.Merge.Internal Data.IntMap.Merge.Lazy Data.IntMap.Merge.Strict Data.IntMap.Strict - Data.IntMap.Strict.Internal Data.IntSet Data.IntSet.Internal Data.Map @@ -90,6 +90,7 @@ library Data.IntMap.Internal.DeprecatedDebug Data.Map.Internal.DeprecatedShowTree Utils.Containers.Internal.TypeError + Utils.Containers.Internal.IsList if impl(ghc >= 8.6) ghc-options: -Werror @@ -319,14 +320,13 @@ test-suite intmap-lazy-properties hs-source-dirs: tests main-is: intmap-properties.hs type: exitcode-stdio-1.0 - other-modules: IntMapValidity build-depends: containers-tests build-depends: array >=0.4.0.0 , base >=4.6 && <5 , deepseq >=1.2 && <1.5 - ghc-options: -O2 + ghc-options: -O1 other-extensions: BangPatterns CPP @@ -343,14 +343,13 @@ test-suite intmap-strict-properties main-is: intmap-properties.hs type: exitcode-stdio-1.0 cpp-options: -DSTRICT - other-modules: IntMapValidity build-depends: containers-tests build-depends: array >=0.4.0.0 , base >=4.6 && <5 , deepseq >=1.2 && <1.5 - ghc-options: -O2 + ghc-options: -O1 other-extensions: BangPatterns CPP diff --git a/containers-tests/tests/IntMapValidity.hs b/containers-tests/tests/IntMapValidity.hs deleted file mode 100644 index 9e92ba631..000000000 --- a/containers-tests/tests/IntMapValidity.hs +++ /dev/null @@ -1,65 +0,0 @@ -module IntMapValidity (valid) where - -import Data.Bits (xor, (.&.)) -import Data.IntMap.Internal -import Test.QuickCheck (Property, counterexample, property, (.&&.)) -import Utils.Containers.Internal.BitUtil (bitcount) - -{-------------------------------------------------------------------- - Assertions ---------------------------------------------------------------------} --- | Returns true iff the internal structure of the IntMap is valid. -valid :: IntMap a -> Property -valid t = - counterexample "nilNeverChildOfBin" (nilNeverChildOfBin t) .&&. - counterexample "commonPrefix" (commonPrefix t) .&&. - counterexample "maskRespected" (maskRespected t) - --- Invariant: Nil is never found as a child of Bin. -nilNeverChildOfBin :: IntMap a -> Bool -nilNeverChildOfBin t = - case t of - Nil -> True - Tip _ _ -> True - Bin _ _ l r -> noNilInSet l && noNilInSet r - where - noNilInSet t' = - case t' of - Nil -> False - Tip _ _ -> True - Bin _ _ l' r' -> noNilInSet l' && noNilInSet r' - --- Invariant: The Mask is a power of 2. It is the largest bit position at which --- two keys of the map differ. -maskPowerOfTwo :: IntMap a -> Bool -maskPowerOfTwo t = - case t of - Nil -> True - Tip _ _ -> True - Bin _ m l r -> - bitcount 0 (fromIntegral m) == 1 && maskPowerOfTwo l && maskPowerOfTwo r - --- Invariant: Prefix is the common high-order bits that all elements share to --- the left of the Mask bit. -commonPrefix :: IntMap a -> Bool -commonPrefix t = - case t of - Nil -> True - Tip _ _ -> True - b@(Bin p _ l r) -> all (sharedPrefix p) (keys b) && commonPrefix l && commonPrefix r - where - sharedPrefix :: Prefix -> Int -> Bool - sharedPrefix p a = p == p .&. a - --- Invariant: In Bin prefix mask left right, left consists of the elements that --- don't have the mask bit set; right is all the elements that do. -maskRespected :: IntMap a -> Bool -maskRespected t = - case t of - Nil -> True - Tip _ _ -> True - Bin _ binMask l r -> - all (\x -> zero x binMask) (keys l) && - all (\x -> not (zero x binMask)) (keys r) && - maskRespected l && - maskRespected r diff --git a/containers-tests/tests/intmap-properties.hs b/containers-tests/tests/intmap-properties.hs index 4a55bc0b9..1d409e5e1 100644 --- a/containers-tests/tests/intmap-properties.hs +++ b/containers-tests/tests/intmap-properties.hs @@ -2,28 +2,32 @@ #ifdef STRICT import Data.IntMap.Strict as Data.IntMap hiding (showTree) -import Data.IntMap.Strict.Internal (traverseMaybeWithKey) import Data.IntMap.Merge.Strict #else import Data.IntMap.Lazy as Data.IntMap hiding (showTree) -import Data.IntMap.Internal (traverseMaybeWithKey) import Data.IntMap.Merge.Lazy #endif -import Data.IntMap.Internal.Debug (showTree) -import IntMapValidity (valid) +import Data.IntMap.Merge.Internal (runWhenMissingAll) +import Data.IntMap.Internal.Debug (showTree, valid, validWith) import Control.Applicative (Applicative(..)) import Data.Monoid import Data.Maybe hiding (mapMaybe) import qualified Data.Maybe as Maybe (mapMaybe) import Data.Ord -import Data.Foldable (foldMap) +import qualified Data.Foldable as Foldable import Data.Function import Data.Traversable (Traversable(traverse), foldMapDefault) +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity(..)) +#else +import Data.IntMap.Merge.Internal (Identity(..)) +#endif import Prelude hiding (lookup, null, map, filter, foldr, foldl) import qualified Prelude (map) import Data.List (nub,sort) +import qualified Data.Foldable as Foldable import qualified Data.List as List import qualified Data.IntSet as IntSet import Test.Framework @@ -141,19 +145,35 @@ main = defaultMain , testProperty "insert to singleton" prop_singleton , testProperty "insert then lookup" prop_insertLookup , testProperty "insert then delete" prop_insertDelete + , testProperty "insertLookupWithKey" prop_insertLookupWithKeyModel , testProperty "delete non member" prop_deleteNonMember + , testProperty "delete" prop_deleteModel + , testProperty "adjust" prop_adjustModel + , testProperty "adjustWithKey" prop_adjustWithKeyModel + , testProperty "update" prop_updateModel + , testProperty "updateWithKey" prop_updateWithKeyModel + , testProperty "updateLookupWithKey" prop_updateLookupWithKeyModel + , testProperty "alter" prop_alterModel + , testProperty "alter lookup" prop_alterLookup + , testProperty "alterF" prop_alterFModel + , testProperty "alterF lookup" prop_alterFLookup , testProperty "union model" prop_unionModel , testProperty "union singleton" prop_unionSingleton , testProperty "union associative" prop_unionAssoc , testProperty "union+unionWith" prop_unionWith , testProperty "union sum" prop_unionSum , testProperty "difference model" prop_differenceModel + , testProperty "differenceWithKey model" prop_differenceWithKeyModel , testProperty "intersection model" prop_intersectionModel , testProperty "intersectionWith model" prop_intersectionWithModel , testProperty "intersectionWithKey model" prop_intersectionWithKeyModel , testProperty "mergeWithKey model" prop_mergeWithKeyModel , testProperty "merge valid" prop_merge_valid , testProperty "mergeA effects" prop_mergeA_effects + , testProperty "union==merge" prop_unionEqMerge + , testProperty "difference==merge" prop_differenceEqMerge + , testProperty "intersection==merge" prop_intersectionEqMerge + , testProperty "merge==mergeA" prop_mergeEqMergeA , testProperty "fromAscList" prop_ordered , testProperty "fromList then toList" prop_list , testProperty "toDescList" prop_descList @@ -182,11 +202,21 @@ main = defaultMain , testProperty "deleteMax" prop_deleteMaxModel , testProperty "filter" prop_filter , testProperty "partition" prop_partition + , testProperty "partitionWithKey" prop_partitionWithKey + , testProperty "mapMaybe" prop_mapMaybeModel + , testProperty "mapMaybeWithKey" prop_mapMaybeWithKeyModel + , testProperty "mapEither" prop_mapEitherModel + , testProperty "mapEitherWithKey" prop_mapEitherWithKeyModel , testProperty "map" prop_map , testProperty "fmap" prop_fmap , testProperty "mapkeys" prop_mapkeys , testProperty "split" prop_splitModel + , testProperty "splitLookup" prop_splitLookupModel , testProperty "splitRoot" prop_splitRoot + , testProperty "isSubmapOf" prop_isSubmapOf + , testProperty "isSubmapOfBy" prop_isSubmapOfBy + , testProperty "isProperSubmapOf" prop_isProperSubmapOf + , testProperty "isProperSubmapOfBy" prop_isProperSubmapOfBy , testProperty "foldr" prop_foldr , testProperty "foldr'" prop_foldr' , testProperty "foldl" prop_foldl @@ -198,6 +228,9 @@ main = defaultMain , testProperty "prop_FoldableTraversableCompat" prop_FoldableTraversableCompat +#if MIN_VERSION_base(4,8,0) + , testProperty "elem" prop_elem +#endif , testProperty "keysSet" prop_keysSet , testProperty "fromSet" prop_fromSet , testProperty "restrictKeys" prop_restrictKeys @@ -207,15 +240,12 @@ main = defaultMain , testProperty "traverseMaybeWithKey identity" prop_traverseMaybeWithKey_identity , testProperty "traverseMaybeWithKey->mapMaybeWithKey" prop_traverseMaybeWithKey_degrade_to_mapMaybeWithKey , testProperty "traverseMaybeWithKey->traverseWithKey" prop_traverseMaybeWithKey_degrade_to_traverseWithKey + , testProperty "filterMissing==filterWithKey" prop_filterMissingEqFilterWithKey + , testProperty "filterAMissing->filterMissing" prop_filterAMissing_degrade_to_filterMissing + , testProperty "mapMissing==mapWithKey" prop_mapMissingEqMapWithKey + , testProperty "traverseMissing->mapMissing" prop_traverseMissing_degrade_to_mapMissing ] -apply2 :: Fun (a, b) c -> a -> b -> c -apply2 f a b = apply f (a, b) - -apply3 :: Fun (a, b, c) d -> a -> b -> c -> d -apply3 f a b c = apply f (a, b, c) - - {-------------------------------------------------------------------- Arbitrary, reasonably balanced trees --------------------------------------------------------------------} @@ -228,6 +258,15 @@ newtype NonEmptyIntMap a = NonEmptyIntMap {getNonEmptyIntMap :: IntMap a} derivi instance Arbitrary a => Arbitrary (NonEmptyIntMap a) where arbitrary = fmap (NonEmptyIntMap . fromList . getNonEmpty) arbitrary +-- | A wrapper around IntMap with a Show instance based on showTree to aid in debugging when +-- tests fail +newtype PrettyIntMap a = PIM { unPIM :: IntMap a } deriving (Eq) + +instance Arbitrary a => Arbitrary (PrettyIntMap a) where + arbitrary = fmap PIM arbitrary + +instance Show a => Show (PrettyIntMap a) where + show (PIM m) = (if valid m then "\n" else "\nINVALID:\n") ++ showTree m ------------------------------------------------------------------------ @@ -237,13 +276,18 @@ type SMap = IntMap String ---------------------------------------------------------------- -tests :: [Test] -tests = [ testGroup "Test Case" [ - ] - , testGroup "Property Test" [ - ] - ] +-- | Like @'nub' . 'sort'@, but more efficient +sortNub :: Ord a => [a] -> [a] +sortNub = sortNubBy compare +sortNubBy :: (a -> a -> Ordering) -> [a] -> [a] +sortNubBy comp = fmap List.head . List.groupBy (\x y -> comp x y == EQ) . List.sortBy comp + +validProp :: IntMap a -> Property +validProp = validWith (flip counterexample) (.&&.) + +allProp :: Testable prop => (a -> prop) -> [a] -> Property +allProp f xs = conjoin (fmap f xs) ---------------------------------------------------------------- -- Unit tests @@ -260,6 +304,8 @@ test_index = do test_index_lookup :: Assertion test_index_lookup = do + (empty :: SMap) !? 1 @?= Nothing + fromList [(5,'a'), (3,'b')] !? 1 @?= Nothing fromList [(5,'a'), (3,'b')] !? 5 @?= Just 'a' @@ -287,6 +333,8 @@ test_size2 = do test_member :: Assertion test_member = do + member 5 empty @?= False + member 5 (fromList [(5,'a'), (3,'b')]) @?= True member 1 (fromList [(5,'a'), (3,'b')]) @?= False @@ -296,6 +344,8 @@ test_member = do test_notMember :: Assertion test_notMember = do + notMember 5 empty @?= True + notMember 5 (fromList [(5,'a'), (3,'b')]) @?= False notMember 1 (fromList [(5,'a'), (3,'b')]) @?= True @@ -319,6 +369,8 @@ test_lookup = do test_findWithDefault :: Assertion test_findWithDefault = do + findWithDefault 'x' 1 empty @?= 'x' + findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) @?= 'x' findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) @?= 'a' @@ -330,6 +382,8 @@ test_findWithDefault = do test_lookupLT :: Assertion test_lookupLT = do + lookupLT 3 (empty :: SMap) @?= Nothing + lookupLT 3 (fromList [(3,'a'), (5,'b')]) @?= Nothing lookupLT 4 (fromList [(3,'a'), (5,'b')]) @?= Just (3, 'a') @@ -340,6 +394,8 @@ test_lookupLT = do test_lookupGT :: Assertion test_lookupGT = do + lookupGT 4 (empty :: SMap) @?= Nothing + lookupGT 4 (fromList [(3,'a'), (5,'b')]) @?= Just (5, 'b') lookupGT 5 (fromList [(3,'a'), (5,'b')]) @?= Nothing @@ -350,6 +406,8 @@ test_lookupGT = do test_lookupLE :: Assertion test_lookupLE = do + lookupLE 2 (empty :: SMap) @?= Nothing + 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') @@ -362,6 +420,8 @@ test_lookupLE = do test_lookupGE :: Assertion test_lookupGE = do + lookupGE 3 (empty :: SMap) @?= Nothing + 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 @@ -489,6 +549,8 @@ test_adjustWithKey = do test_update :: Assertion test_update = do + update f 5 empty @?= empty + 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" @@ -502,6 +564,8 @@ test_update = do test_updateWithKey :: Assertion test_updateWithKey = do + updateWithKey f 5 empty @?= empty + 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" @@ -515,6 +579,8 @@ test_updateWithKey = do test_updateLookupWithKey :: Assertion test_updateLookupWithKey = do + updateLookupWithKey f 5 empty @?= (Nothing, empty) + 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") @@ -528,6 +594,9 @@ test_updateLookupWithKey = do test_alter :: Assertion test_alter = do + alter f 7 empty @?= (empty :: SMap) + alter g 7 empty @?= singleton 7 "c" + alter f 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a")] alter f 5 (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "b" alter g 7 (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "b"), (5, "a"), (7, "c")] @@ -550,21 +619,25 @@ test_alter = do test_union :: Assertion test_union = do + union empty empty @?= (empty :: SMap) union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "a"), (7, "C")] union (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(-3, "b"), (5, "a"), (7, "C")] test_mappend :: Assertion test_mappend = do + mappend empty empty @?= (empty :: SMap) mappend (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "a"), (7, "C")] mappend (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(-3, "b"), (5, "a"), (7, "C")] test_unionWith :: Assertion test_unionWith = do + unionWith (++) empty empty @?= (empty :: SMap) unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "aA"), (7, "C")] unionWith (++) (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(-3, "b"), (5, "aA"), (7, "C")] test_unionWithKey :: Assertion test_unionWithKey = do + unionWithKey f empty empty @?= (empty :: SMap) unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(3, "b"), (5, "5:a|A"), (7, "C")] unionWithKey f (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= fromList [(-3, "b"), (5, "5:a|A"), (7, "C")] where @@ -572,6 +645,8 @@ test_unionWithKey = do test_unions :: Assertion test_unions = do + unions [] @?= (empty :: SMap) + 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")])] @@ -584,6 +659,8 @@ test_unions = do test_mconcat :: Assertion test_mconcat = do + mconcat [] @?= (empty :: SMap) + mconcat [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] @?= fromList [(3, "b"), (5, "a"), (7, "C")] mconcat [(fromList [(5, "A3"), (3, "B3")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "a"), (3, "b")])] @@ -596,6 +673,7 @@ test_mconcat = do test_unionsWith :: Assertion test_unionsWith = do + unionsWith (++) [] @?= (empty :: SMap) unionsWith (++) [(fromList [(5, "a"), (3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (3, "B3")])] @?= fromList [(3, "bB3"), (5, "aAA3"), (7, "C")] unionsWith (++) [(fromList [(5, "a"), (-3, "b")]), (fromList [(5, "A"), (7, "C")]), (fromList [(5, "A3"), (-3, "B3")])] @@ -603,11 +681,17 @@ test_unionsWith = do test_difference :: Assertion test_difference = do + difference empty empty @?= (empty :: SMap) + difference empty (fromList [(5, "A"), (7, "C")]) @?= (empty :: SMap) + difference (fromList [(5, "a"), (3, "b")]) empty @?= fromList [(5, "a"), (3, "b")] difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 3 "b" difference (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton (-3) "b" test_differenceWith :: Assertion test_differenceWith = do + differenceWith f empty empty @?= (empty :: SMap) + differenceWith f empty (fromList [(5, "A"), (3, "B"), (7, "C")]) @?= (empty :: SMap) + differenceWith f (fromList [(5, "a"), (3, "b")]) empty @?= fromList [(5, "a"), (3, "b")] differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) @?= singleton 3 "b:B" differenceWith f (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (-3, "B"), (7, "C")]) @@ -617,6 +701,9 @@ test_differenceWith = do test_differenceWithKey :: Assertion test_differenceWithKey = do + differenceWithKey f empty empty @?= (empty :: SMap) + differenceWithKey f empty (fromList [(5, "A"), (3, "B"), (7, "C")]) @?= empty + differenceWithKey f (fromList [(5, "a"), (3, "b")]) empty @?= fromList [(5, "a"), (3, "b")] differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) @?= singleton 3 "3:b|B" differenceWithKey f (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (-3, "B"), (10, "C")]) @@ -626,17 +713,26 @@ test_differenceWithKey = do test_intersection :: Assertion test_intersection = do + intersection empty empty @?= (empty :: SMap) + intersection empty (fromList [(5, "A"), (7, "C")]) @?= (empty :: SMap) + intersection (fromList [(5, "a"), (3, "b")]) empty @?= empty intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "a" intersection (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "a" test_intersectionWith :: Assertion test_intersectionWith = do + intersectionWith (++) empty empty @?= (empty :: SMap) + intersectionWith (++) empty (fromList [(5, "A"), (7, "C")]) @?= empty + intersectionWith (++) (fromList [(5, "a"), (3, "b")]) empty @?= empty intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "aA" intersectionWith (++) (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "aA" test_intersectionWithKey :: Assertion test_intersectionWithKey = do + intersectionWithKey f empty empty @?= (empty :: SMap) + intersectionWithKey f empty (fromList [(5, "A"), (7, "C")]) @?= empty + intersectionWithKey f (fromList [(5, "a"), (3, "b")]) empty @?= empty intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "5:a|A" intersectionWithKey f (fromList [(5, "a"), (-3, "b")]) (fromList [(5, "A"), (7, "C")]) @?= singleton 5 "5:a|A" where @@ -647,12 +743,14 @@ test_intersectionWithKey = do test_map :: Assertion test_map = do + map (++ "x") empty @?= empty map (++ "x") (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "bx"), (5, "ax")] map (++ "x") (fromList [(5,"a"), (3,"b"), (-1,"c")]) @?= fromList [(3, "bx"), (5, "ax"), (-1,"cx")] test_mapWithKey :: Assertion test_mapWithKey = do + mapWithKey f empty @?= empty mapWithKey f (fromList [(5,"a"), (3,"b")]) @?= fromList [(3, "3:b"), (5, "5:a")] mapWithKey f (fromList [(5,"a"), (3,"b"), (-1,"c")]) @?= fromList [(3, "3:b"), (5, "5:a"), (-1,"-1:c")] @@ -661,6 +759,7 @@ test_mapWithKey = do test_mapAccum :: Assertion test_mapAccum = do + mapAccum f "Everything: " empty @?= ("Everything: ", empty) mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) @?= ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) mapAccum f "Everything: " (fromList [(5,"a"), (3,"b"), (-1,"c")]) @?= ("Everything: cba", fromList [(3, "bX"), (5, "aX"), (-1, "cX")]) @@ -669,6 +768,7 @@ test_mapAccum = do test_mapAccumWithKey :: Assertion test_mapAccumWithKey = do + mapAccumWithKey f "Everything:" empty @?= ("Everything:", empty) mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) @?= ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b"), (-1,"c")]) @?= ("Everything: -1-c 3-b 5-a", fromList [(3, "bX"), (5, "aX"), (-1,"cX")]) @@ -677,6 +777,7 @@ test_mapAccumWithKey = do test_mapAccumRWithKey :: Assertion test_mapAccumRWithKey = do + mapAccumRWithKey f "Everything:" empty @?= ("Everything:", empty) mapAccumRWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) @?= ("Everything: 5-a 3-b", fromList [(3, "bX"), (5, "aX")]) mapAccumRWithKey f "Everything:" (fromList [(5,"a"), (3,"b"), (-1,"c")]) @?= ("Everything: 5-a 3-b -1-c", fromList [(3, "bX"), (5, "aX"), (-1,"cX")]) @@ -685,6 +786,8 @@ test_mapAccumRWithKey = do test_mapKeys :: Assertion test_mapKeys = do + mapKeys (+ 1) empty @?= (empty :: SMap) + 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" @@ -698,6 +801,8 @@ test_mapKeys = do test_mapKeysWith :: Assertion test_mapKeysWith = do + mapKeysWith (++) (\ _ -> 1) empty @?= (empty :: SMap) + 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" @@ -706,6 +811,8 @@ test_mapKeysWith = do test_mapKeysMonotonic :: Assertion test_mapKeysMonotonic = do + mapKeysMonotonic (+ 1) empty @?= (empty :: SMap) + mapKeysMonotonic (+ 1) (fromList [(5,"a"), (3,"b")]) @?= fromList [(4, "b"), (6, "a")] mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) @?= fromList [(6, "b"), (10, "a")] @@ -784,46 +891,53 @@ test_fromListWithKey = do test_toAscList :: Assertion test_toAscList = do + toAscList (empty :: SMap) @?= [] toAscList (fromList [(5,"a"), (3,"b")]) @?= [(3,"b"), (5,"a")] toAscList (fromList [(5,"a"), (-3,"b")]) @?= [(-3,"b"), (5,"a")] test_toDescList :: Assertion test_toDescList = do + toDescList (empty :: SMap) @?= [] toDescList (fromList [(5,"a"), (3,"b")]) @?= [(5,"a"), (3,"b")] toDescList (fromList [(5,"a"), (-3,"b")]) @?= [(5,"a"), (-3,"b")] test_showTree :: Assertion test_showTree = do + showTree (empty :: UMap) @?= unlines [] showTree posTree @?= expectedPosTree showTree negTree @?= expectedNegTree where mkAscTree ls = fromDistinctAscList [(x,()) | x <- ls] posTree = mkAscTree [1..5] negTree = mkAscTree [(-2)..2] expectedPosTree = unlines - [ "*" - , "+--*" - , "| +-- 1:=()" - , "| +--*" - , "| +-- 2:=()" - , "| +-- 3:=()" - , "+--*" - , " +-- 4:=()" - , " +-- 5:=()" + [ "1:=()" + , "| ,-*" + , "| +---. 2:=()" + , "| | +-*" + , "| | `-*" + , "| ,---' 3:=()" + , "| +---. 4:=()" + , "| | +-*" + , "| | `-*" + , "`---' 5:=()" ] expectedNegTree = unlines - [ "*" - , "+--*" - , "| +--*" - , "| | +-- 0:=()" - , "| | +-- 1:=()" - , "| +-- 2:=()" - , "+--*" - , " +-- -2:=()" - , " +-- -1:=()" + [ "-2:=()" + , "| ,-*" + , "| +-*" + , "| ,---' -1:=()" + , "| +---. 0:=()" + , "| | | ,-*" + , "| | | +-*" + , "| | +---' 1:=()" + , "| | `-*" + , "`---' 2:=()" ] test_fromAscList :: Assertion test_fromAscList = do + fromAscList [] @?= (empty :: SMap) + fromAscList [(3,"b"), (5,"a")] @?= fromList [(3, "b"), (5, "a")] fromAscList [(3,"b"), (5,"a"), (5,"b")] @?= fromList [(3, "b"), (5, "b")] @@ -833,11 +947,15 @@ test_fromAscList = do test_fromAscListWith :: Assertion test_fromAscListWith = do + fromAscListWith (++) [] @?= (empty :: SMap) + fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] @?= fromList [(3, "b"), (5, "ba")] fromAscListWith (++) [(-3,"b"), (5,"a"), (5,"b")] @?= fromList [(-3, "b"), (5, "ba")] test_fromAscListWithKey :: Assertion test_fromAscListWithKey = do + fromAscListWithKey f [] @?= (empty :: SMap) + fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b"), (5,"b")] @?= fromList [(3, "b"), (5, "5:b5:ba")] fromAscListWithKey f [(-3,"b"), (5,"a"), (5,"b"), (5,"b")] @?= fromList [(-3, "b"), (5, "5:b5:ba")] where @@ -845,6 +963,7 @@ test_fromAscListWithKey = do test_fromDistinctAscList :: Assertion test_fromDistinctAscList = do + fromDistinctAscList [] @?= (empty :: SMap) fromDistinctAscList [(3,"b"), (5,"a")] @?= fromList [(3, "b"), (5, "a")] fromDistinctAscList [(-3,"b"), (5,"a")] @?= fromList [(-3, "b"), (5, "a")] @@ -853,6 +972,8 @@ test_fromDistinctAscList = do test_filter :: Assertion test_filter = do + filter (> "a") (empty :: SMap) @?= empty + 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 @@ -863,11 +984,15 @@ test_filter = do test_filteWithKey :: Assertion test_filteWithKey = do + filterWithKey (\k _ -> k > 4) (empty :: SMap) @?= empty + filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "a" filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (-3,"b")]) @?= singleton 5 "a" test_partition :: Assertion test_partition = do + partition (> "a") (empty :: SMap) @?= (empty, empty) + 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")]) @@ -878,6 +1003,8 @@ test_partition = do test_partitionWithKey :: Assertion test_partitionWithKey = do + partitionWithKey (\ k _ -> k > 3) (empty :: SMap) @?= (empty, empty) + 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")]) @@ -888,6 +1015,7 @@ test_partitionWithKey = do test_mapMaybe :: Assertion test_mapMaybe = do + mapMaybe f (empty :: SMap) @?= empty mapMaybe f (fromList [(5,"a"), (3,"b")]) @?= singleton 5 "new a" mapMaybe f (fromList [(5,"a"), (-3,"b")]) @?= singleton 5 "new a" where @@ -895,6 +1023,7 @@ test_mapMaybe = do test_mapMaybeWithKey :: Assertion test_mapMaybeWithKey = do + mapMaybeWithKey f (empty :: SMap) @?= empty mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) @?= singleton 3 "key : 3" mapMaybeWithKey f (fromList [(5,"a"), (-3,"b")]) @?= singleton (-3) "key : -3" where @@ -902,6 +1031,8 @@ test_mapMaybeWithKey = do test_mapEither :: Assertion test_mapEither = do + mapEither f (empty :: SMap) @?= (empty, empty) + 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")]) @@ -916,6 +1047,8 @@ test_mapEither = do test_mapEitherWithKey :: Assertion test_mapEitherWithKey = do + mapEitherWithKey f (empty :: SMap) @?= (empty, empty) + 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")]) @@ -930,6 +1063,8 @@ test_mapEitherWithKey = do test_split :: Assertion test_split = do + split 2 (empty :: SMap) @?= (empty, empty) + 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") @@ -944,6 +1079,8 @@ test_split = do test_splitLookup :: Assertion test_splitLookup = do + splitLookup 2 (empty :: SMap) @?= (empty, Nothing, empty) + 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") @@ -961,6 +1098,10 @@ test_splitLookup = do test_isSubmapOfBy :: Assertion test_isSubmapOfBy = do + isSubmapOfBy (==) empty (empty :: IMap) @?= True + isSubmapOfBy (==) (fromList [(-1,1),(2,2)]) empty @?= False + isSubmapOfBy (==) empty (fromList [(-1,1),(2,2)]) @?= True + isSubmapOfBy (==) (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True isSubmapOfBy (<=) (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True isSubmapOfBy (==) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True @@ -978,6 +1119,10 @@ test_isSubmapOfBy = do test_isSubmapOf :: Assertion test_isSubmapOf = do + isSubmapOf empty (empty :: IMap) @?= True + isSubmapOf (fromList [(-1,1),(2,2)]) empty @?= False + isSubmapOf empty (fromList [(-1,1),(2,2)]) @?= True + isSubmapOf (fromList [(fromEnum 'a',1)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True isSubmapOf (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= True isSubmapOf (fromList [(fromEnum 'a',2)]) (fromList [(fromEnum 'a',1),(fromEnum 'b',2)]) @?= False @@ -990,6 +1135,10 @@ test_isSubmapOf = do test_isProperSubmapOfBy :: Assertion test_isProperSubmapOfBy = do + isProperSubmapOfBy (==) empty (empty :: IMap) @?= False + isProperSubmapOfBy (==) (fromList [(-1,1),(2,2)]) empty @?= False + isProperSubmapOfBy (==) empty (fromList [(-1,1),(2,2)]) @?= True + isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True isProperSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) @?= False @@ -1004,6 +1153,10 @@ test_isProperSubmapOfBy = do test_isProperSubmapOf :: Assertion test_isProperSubmapOf = do + isProperSubmapOf empty (empty :: IMap) @?= False + isProperSubmapOf (fromList [(-1,1),(2,2)]) empty @?= False + isProperSubmapOf empty (fromList [(-1,1),(2,2)]) @?= True + isProperSubmapOf (fromList [(1,1)]) (fromList [(1,1),(2,2)]) @?= True isProperSubmapOf (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) @?= False isProperSubmapOf (fromList [(1,1),(2,2)]) (fromList [(1,1)]) @?= False @@ -1155,20 +1308,20 @@ forValidUnitTree :: Testable b => (SMap -> b) -> Property forValidUnitTree f = forValid f prop_valid :: Property -prop_valid = forValidUnitTree $ \t -> valid t +prop_valid = forValidUnitTree $ \t -> validProp t ---------------------------------------------------------------- -- QuickCheck ---------------------------------------------------------------- prop_emptyValid :: Property -prop_emptyValid = valid empty +prop_emptyValid = validProp empty prop_singleton :: Int -> Int -> Property prop_singleton k x = case singleton k x of s -> - valid s .&&. + validProp s .&&. s === insert k x empty prop_insertLookup :: Int -> UMap -> Bool @@ -1178,66 +1331,136 @@ prop_insertDelete :: Int -> UMap -> Property prop_insertDelete k t = lookup k t == Nothing ==> case delete k (insert k () t) of - t' -> valid t' .&&. t' === t + t' -> validProp t' .&&. t' === t + +prop_insertLookupWithKeyModel :: Fun (Key, A, A) A -> Key -> A -> IntMap A -> Property +prop_insertLookupWithKeyModel fun k v m = + insertLookupWithKey (applyFun3 fun) k v m + === (lookup k m, insertWithKey (applyFun3 fun) k v m) prop_deleteNonMember :: Int -> UMap -> Property -prop_deleteNonMember k t = (lookup k t == Nothing) ==> (delete k t == t) +prop_deleteNonMember k t = notMember k t ==> (delete k t === t) ---------------------------------------------------------------- +prop_deleteModel :: Key -> IntMap A -> Property +prop_deleteModel k m = + delete k m === fromList (List.filter ((/= k) . fst) (toList m)) + +prop_adjustModel :: Fun A A -> Key -> IntMap A -> Property +prop_adjustModel fun k m = + adjust (apply fun) k m === alter (fmap (apply fun)) k m + +prop_adjustWithKeyModel :: Fun (Key, A) A -> Key -> IntMap A -> Property +prop_adjustWithKeyModel fun k m = + adjustWithKey (applyFun2 fun) k m === adjust (applyFun2 fun k) k m + +prop_updateModel :: Fun A (Maybe A) -> Key -> IntMap A -> Property +prop_updateModel fun k m = + update (apply fun) k m === alter (>>= apply fun) k m + +prop_updateWithKeyModel :: Fun (Key, A) (Maybe A) -> Key -> IntMap A -> Property +prop_updateWithKeyModel fun k m = + updateWithKey (applyFun2 fun) k m === update (applyFun2 fun k) k m + +prop_updateLookupWithKeyModel :: Fun (Key, A) (Maybe A) -> Key -> IntMap A -> Property +prop_updateLookupWithKeyModel fun k m = + updateLookupWithKey (applyFun2 fun) k m + === (lookup k m, updateWithKey (applyFun2 fun) k m) + +prop_alterModel :: Fun (Maybe A) (Maybe A) -> Key -> IntMap A -> Property +prop_alterModel fun k m = + alter (apply fun) k m + === let old = lookup k m + in case (old, apply fun old) of + (Nothing, Nothing) -> m + (Just _, Nothing) -> delete k m + (_, Just v) -> insert k v m + +prop_alterLookup :: Fun (Maybe A) (Maybe A) -> Key -> IntMap A -> Property +prop_alterLookup fun k m = + lookup k (alter (apply fun) k m) === apply fun (lookup k m) + +prop_alterFModel :: Fun (Maybe A) (Maybe A) -> Key -> IntMap A -> Property +prop_alterFModel fun k m = + alterF f k m + === let old = lookup k m + in (old, case (old, apply fun old) of + (Nothing, Nothing) -> m + (Just _, Nothing) -> delete k m + (_, Just v) -> insert k v m) + where + f mv = (mv, apply fun mv) + +prop_alterFLookup :: Fun (Maybe A) (Maybe A) -> Key -> IntMap A -> Property +prop_alterFLookup fun k m = + fmap (lookup k) (alterF f k m) === f (lookup k m) + where + f mv = (mv, apply fun mv) + + prop_unionModel :: [(Int,Int)] -> [(Int,Int)] -> Property prop_unionModel xs ys = case union (fromList xs) (fromList ys) of t -> - valid t .&&. - sort (keys t) === sort (nub (Prelude.map fst xs ++ Prelude.map fst ys)) + validProp t .&&. + keys t === sortNub (Prelude.map fst xs ++ Prelude.map fst ys) -prop_unionSingleton :: IMap -> Int -> Int -> Bool -prop_unionSingleton t k x = union (singleton k x) t == insert k x t +prop_unionSingleton :: IMap -> Int -> Int -> Property +prop_unionSingleton t k x = union (singleton k x) t === insert k x t -prop_unionAssoc :: IMap -> IMap -> IMap -> Bool -prop_unionAssoc t1 t2 t3 = union t1 (union t2 t3) == union (union t1 t2) t3 +prop_unionAssoc :: IMap -> IMap -> IMap -> Property +prop_unionAssoc t1 t2 t3 = union t1 (union t2 t3) === union (union t1 t2) t3 -prop_unionWith :: IMap -> IMap -> Bool -prop_unionWith t1 t2 = (union t1 t2 == unionWith (\_ y -> y) t2 t1) +prop_unionWith :: IMap -> IMap -> Property +prop_unionWith t1 t2 = union t1 t2 === unionWith (\_ y -> y) t2 t1 -prop_unionSum :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_unionSum :: [(Int,Int)] -> [(Int,Int)] -> Property prop_unionSum xs ys - = sum (elems (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys))) - == (sum (Prelude.map snd xs) + sum (Prelude.map snd ys)) + = Foldable.sum (unionWith (+) (fromListWith (+) xs) (fromListWith (+) ys)) + === (sum (Prelude.map snd xs) + sum (Prelude.map snd ys)) prop_differenceModel :: [(Int,Int)] -> [(Int,Int)] -> Property prop_differenceModel xs ys = case difference (fromListWith (+) xs) (fromListWith (+) ys) of t -> - valid t .&&. - sort (keys t) === sort ((List.\\) - (nub (Prelude.map fst xs)) - (nub (Prelude.map fst ys))) + validProp t .&&. + keys t === (List.\\) (sortNub (Prelude.map fst xs)) + (sortNub (Prelude.map fst ys)) + +prop_differenceWithKeyModel :: Fun (Int, Int, Int) (Maybe Int) -> [(Int,Int)] -> [(Int,Int)] -> Property +prop_differenceWithKeyModel f xs ys + = differenceWithKey (\k x y -> apply f (k, x, y)) (fromList xs') (fromList ys') + === fromList (Maybe.mapMaybe diffSingle xs') + where + xs' = sortNubBy (compare `on` fst) xs + ys' = sortNubBy (compare `on` fst) ys + diffSingle (k, x) = case List.lookup k ys' of + Nothing -> Just (k, x) + Just y -> fmap (\r -> (k, r)) (apply f (k, x, y)) prop_intersectionModel :: [(Int,Int)] -> [(Int,Int)] -> Property prop_intersectionModel xs ys = case intersection (fromListWith (+) xs) (fromListWith (+) ys) of t -> - valid t .&&. - sort (keys t) === sort (nub ((List.intersect) - (Prelude.map fst xs) - (Prelude.map fst ys))) + validProp t .&&. + keys t === (List.intersect) (sortNub (Prelude.map fst xs)) + (sortNub (Prelude.map fst ys)) -prop_intersectionWithModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionWithModel :: [(Int,Int)] -> [(Int,Int)] -> Property prop_intersectionWithModel xs ys - = toList (intersectionWith f (fromList xs') (fromList ys')) - == [(kx, f vx vy ) | (kx, vx) <- List.sort xs', (ky, vy) <- ys', kx == ky] - where xs' = List.nubBy ((==) `on` fst) xs - ys' = List.nubBy ((==) `on` fst) ys + = intersectionWith f (fromList xs') (fromList ys') + === fromList [(kx, f vx vy ) | (kx, vx) <- xs', (ky, vy) <- ys', kx == ky] + where xs' = sortNubBy (compare `on` fst) xs + ys' = sortNubBy (compare `on` fst) ys f l r = l + 2 * r -prop_intersectionWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_intersectionWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Property prop_intersectionWithKeyModel xs ys - = toList (intersectionWithKey f (fromList xs') (fromList ys')) - == [(kx, f kx vx vy) | (kx, vx) <- List.sort xs', (ky, vy) <- ys', kx == ky] - where xs' = List.nubBy ((==) `on` fst) xs - ys' = List.nubBy ((==) `on` fst) ys + = intersectionWithKey f (fromList xs') (fromList ys') + === fromList [(kx, f kx vx vy) | (kx, vx) <- xs', (ky, vy) <- ys', kx == ky] + where xs' = sortNubBy (compare `on` fst) xs + ys' = sortNubBy (compare `on` fst) ys f k l r = k + 2 * l + 3 * r prop_disjoint :: UMap -> UMap -> Property @@ -1259,26 +1482,26 @@ prop_withoutKeys m s0 = where s = keysSet s0 -prop_mergeWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Bool +prop_mergeWithKeyModel :: [(Int,Int)] -> [(Int,Int)] -> Property prop_mergeWithKeyModel xs ys - = and [ testMergeWithKey f keep_x keep_y - | f <- [ \_k x1 _x2 -> Just x1 - , \_k _x1 x2 -> Just x2 - , \_k _x1 _x2 -> Nothing - , \k x1 x2 -> if k `mod` 2 == 0 then Nothing else Just (2 * x1 + 3 * x2) - ] - , keep_x <- [ True, False ] - , keep_y <- [ True, False ] - ] - - where xs' = List.nubBy ((==) `on` fst) xs - ys' = List.nubBy ((==) `on` fst) ys + = conjoin [ testMergeWithKey f keep_x keep_y + | f <- [ \_k x1 _x2 -> Just x1 + , \_k _x1 x2 -> Just x2 + , \_k _x1 _x2 -> Nothing + , \k x1 x2 -> if k `mod` 2 == 0 then Nothing else Just (2 * x1 + 3 * x2) + ] + , keep_x <- [ True, False ] + , keep_y <- [ True, False ] + ] + + where xs' = sortNubBy (compare `on` fst) xs + ys' = sortNubBy (compare `on` fst) ys xm = fromList xs' ym = fromList ys' testMergeWithKey f keep_x keep_y - = toList (mergeWithKey f (keep keep_x) (keep keep_y) xm ym) == emulateMergeWithKey f keep_x keep_y + = mergeWithKey f (keep keep_x) (keep keep_y) xm ym === fromList (emulateMergeWithKey f keep_x keep_y) where keep False _ = empty keep True m = m @@ -1302,7 +1525,7 @@ prop_merge_valid -> IntMap B -> Property prop_merge_valid whenMissingA whenMissingB whenMatched xs ys - = valid m + = validProp m where m = merge @@ -1323,28 +1546,46 @@ prop_mergeA_effects xs ys whenMissing = traverseMissing (\k _ -> ([k], ())) whenMatched = zipWithAMatched (\k _ _ -> ([k], ())) +prop_unionEqMerge :: UMap -> UMap -> Property +prop_unionEqMerge m1 m2 = PIM (union m1 m2) === PIM (merge preserveMissing preserveMissing (zipWithMatched (\_ x _ -> x)) m1 m2) + +prop_differenceEqMerge :: UMap -> UMap -> Property +prop_differenceEqMerge m1 m2 = PIM (difference m1 m2) === PIM (merge preserveMissing dropMissing (zipWithMaybeMatched (\_ _ _ -> Nothing)) m1 m2) + +prop_intersectionEqMerge :: UMap -> UMap -> Property +prop_intersectionEqMerge m1 m2 = PIM (intersection m1 m2) === PIM (merge dropMissing dropMissing (zipWithMatched (\_ x _ -> x)) m1 m2) + +prop_mergeEqMergeA :: Fun Int Bool -> Fun Int Bool -> Fun Int Bool -> UMap -> UMap -> Property +prop_mergeEqMergeA pMiss1 pMiss2 pMatch m1 m2 = PIM merged === PIM mergedA where + merged = merge whenMiss1 whenMiss2 whenMatch m1 m2 + mergedA = runIdentity (mergeA whenMiss1 whenMiss2 whenMatch m1 m2) + + whenMiss1 = mapMaybeMissing (\k _ -> if apply pMiss1 k then Just () else Nothing) + whenMiss2 = mapMaybeMissing (\k _ -> if apply pMiss2 k then Just () else Nothing) + whenMatch = zipWithMaybeMatched (\k _ _ -> if apply pMatch k then Just () else Nothing) + ---------------------------------------------------------------- prop_ordered :: Property prop_ordered = forAll (choose (5,100)) $ \n -> let xs = [(x,()) | x <- [0..n::Int]] - in fromAscList xs == fromList xs + in fromAscList xs === fromList xs -prop_list :: [Int] -> Bool -prop_list xs = (sort (nub xs) == [x | (x,()) <- toList (fromList [(x,()) | x <- xs])]) +prop_list :: [Int] -> Property +prop_list xs = sortNub xs === [x | (x,()) <- toList (fromList [(x,()) | x <- xs])] -prop_descList :: [Int] -> Bool -prop_descList xs = (reverse (sort (nub xs)) == [x | (x,()) <- toDescList (fromList [(x,()) | x <- xs])]) +prop_descList :: [Int] -> Property +prop_descList xs = reverse (sortNub xs) === [x | (x,()) <- toDescList (fromList [(x,()) | x <- xs])] -prop_ascDescList :: [Int] -> Bool -prop_ascDescList xs = toAscList m == reverse (toDescList m) +prop_ascDescList :: [Int] -> Property +prop_ascDescList xs = toAscList m === reverse (toDescList m) where m = fromList $ zip xs $ repeat () prop_fromList :: [Int] -> Property prop_fromList xs = case fromList (zip xs xs) of - t -> valid t .&&. + t -> validProp t .&&. t === fromAscList (zip sort_xs sort_xs) .&&. t === fromDistinctAscList (zip nub_sort_xs nub_sort_xs) .&&. t === List.foldr (uncurry insert) empty (zip xs xs) @@ -1354,9 +1595,9 @@ prop_fromList xs ---------------------------------------------------------------- prop_alter :: UMap -> Int -> Property -prop_alter t k = valid t' .&&. case lookup k t of - Just _ -> (size t - 1) == size t' && lookup k t' == Nothing - Nothing -> (size t + 1) == size t' && lookup k t' /= Nothing +prop_alter t k = validProp t' .&&. case lookup k t of + Just _ -> (size t - 1) === size t' .&&. lookup k t' === Nothing + Nothing -> (size t + 1) === size t' .&&. lookup k t' =/= Nothing where t' = alter f k t f Nothing = Just () @@ -1366,14 +1607,14 @@ prop_alter t k = valid t' .&&. case lookup k t of -- Compare against the list model (after nub on keys) prop_index :: [Int] -> Property -prop_index xs = length xs > 0 ==> +prop_index xs = let m = fromList (zip xs xs) - in xs == [ m ! i | i <- xs ] + in xs === [ m ! i | i <- xs ] prop_index_lookup :: [Int] -> Property -prop_index_lookup xs = length xs > 0 ==> +prop_index_lookup xs = let m = fromList (zip xs xs) - in (Prelude.map Just xs) == [ m !? i | i <- xs ] + in (Prelude.map Just xs) === [ m !? i | i <- xs ] prop_null :: IMap -> Bool prop_null m = null m == (size m == 0) @@ -1383,58 +1624,58 @@ prop_size im = sz === foldl' (\i _ -> i + 1) (0 :: Int) im .&&. sz === List.length (toList im) where sz = size im -prop_member :: [Int] -> Int -> Bool +prop_member :: [Int] -> Int -> Property prop_member xs n = let m = fromList (zip xs xs) - in all (\k -> k `member` m == (k `elem` xs)) (n : xs) + in allProp (\k -> k `member` m === (k `elem` xs)) (n : xs) -prop_notmember :: [Int] -> Int -> Bool +prop_notmember :: [Int] -> Int -> Property prop_notmember xs n = let m = fromList (zip xs xs) - in all (\k -> k `notMember` m == (k `notElem` xs)) (n : xs) + in allProp (\k -> k `notMember` m === (k `notElem` xs)) (n : xs) -prop_lookup :: [(Int, Int)] -> Int -> Bool +prop_lookup :: [(Int, Int)] -> Int -> Property prop_lookup xs n = - let xs' = List.nubBy ((==) `on` fst) xs + let xs' = sortNubBy (compare `on` fst) xs m = fromList xs' - in all (\k -> lookup k m == List.lookup k xs') (n : List.map fst xs') + in allProp (\k -> lookup k m === List.lookup k xs') (n : List.map fst xs') -prop_find :: [(Int, Int)] -> Bool +prop_find :: [(Int, Int)] -> Property prop_find xs = - let xs' = List.nubBy ((==) `on` fst) xs + let xs' = sortNubBy (compare `on` fst) xs m = fromList xs' - in all (\(k, v) -> m ! k == v) xs' + in allProp (\(k, v) -> m ! k === v) xs' -prop_findWithDefault :: [(Int, Int)] -> Int -> Int -> Bool +prop_findWithDefault :: [(Int, Int)] -> Int -> Int -> Property prop_findWithDefault xs n x = - let xs' = List.nubBy ((==) `on` fst) xs + let xs' = sortNubBy (compare `on` fst) xs m = fromList xs' - in all (\k -> findWithDefault x k m == maybe x id (List.lookup k xs')) (n : List.map fst xs') + in allProp (\k -> findWithDefault x k m === maybe x id (List.lookup k xs')) (n : List.map fst xs') -test_lookupSomething :: (Int -> IntMap Int -> Maybe (Int, Int)) -> (Int -> Int -> Bool) -> [(Int, Int)] -> Bool +test_lookupSomething :: (Int -> IntMap Int -> Maybe (Int, Int)) -> (Int -> Int -> Bool) -> [(Int, Int)] -> Property test_lookupSomething lookup' cmp xs = - let odd_sorted_xs = filter_odd $ sort $ List.nubBy ((==) `on` fst) xs + let odd_sorted_xs = filter_odd $ sortNubBy (compare `on` fst) xs t = fromList odd_sorted_xs test k = case List.filter ((`cmp` k) . fst) odd_sorted_xs of - [] -> lookup' k t == Nothing - cs | 0 `cmp` 1 -> lookup' k t == Just (last cs) -- we want largest such element - | otherwise -> lookup' k t == Just (head cs) -- we want smallest such element - in all test (List.map fst xs) + [] -> lookup' k t === Nothing + cs | 0 `cmp` 1 -> lookup' k t === Just (last cs) -- we want largest such element + | otherwise -> lookup' k t === Just (head cs) -- we want smallest such element + in allProp test (List.map fst xs) where filter_odd [] = [] filter_odd [_] = [] filter_odd (_ : o : xs) = o : filter_odd xs -prop_lookupLT :: [(Int, Int)] -> Bool +prop_lookupLT :: [(Int, Int)] -> Property prop_lookupLT = test_lookupSomething lookupLT (<) -prop_lookupGT :: [(Int, Int)] -> Bool +prop_lookupGT :: [(Int, Int)] -> Property prop_lookupGT = test_lookupSomething lookupGT (>) -prop_lookupLE :: [(Int, Int)] -> Bool +prop_lookupLE :: [(Int, Int)] -> Property prop_lookupLE = test_lookupSomething lookupLE (<=) -prop_lookupGE :: [(Int, Int)] -> Bool +prop_lookupGE :: [(Int, Int)] -> Property prop_lookupGE = test_lookupSomething lookupGE (>=) prop_lookupMin :: IntMap Int -> Property @@ -1451,61 +1692,106 @@ prop_findMax (NonEmptyIntMap im) = findMax im === head (toDescList im) prop_deleteMinModel :: [(Int, Int)] -> Property prop_deleteMinModel ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in toAscList (deleteMin m) == tail (sort xs) + in deleteMin m === fromList (tail xs) prop_deleteMaxModel :: [(Int, Int)] -> Property prop_deleteMaxModel ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in toAscList (deleteMax m) == init (sort xs) + in deleteMax m === fromList (init xs) prop_filter :: Fun Int Bool -> [(Int, Int)] -> Property -prop_filter p ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_filter p ys = + let xs = sortNubBy (compare `on` fst) ys m = filter (apply p) (fromList xs) - in valid m .&&. + in validProp m .&&. m === fromList (List.filter (apply p . snd) xs) prop_partition :: Fun Int Bool -> [(Int, Int)] -> Property -prop_partition p ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_partition p ys = + let xs = sortNubBy (compare `on` fst) ys m@(l, r) = partition (apply p) (fromList xs) - in valid l .&&. - valid r .&&. + in validProp l .&&. + validProp r .&&. m === let (a,b) = (List.partition (apply p . snd) xs) in (fromList a, fromList b) +prop_partitionWithKey :: Fun (Int, Int) Bool -> [(Int, Int)] -> Property +prop_partitionWithKey p ys = + let xs = sortNubBy (compare `on` fst) ys + m@(l, r) = partitionWithKey (curry (apply p)) (fromList xs) + in validProp l .&&. + validProp r .&&. + m === let (a,b) = (List.partition (apply p) xs) + in (fromList a, fromList b) + +prop_mapMaybeModel :: Fun A (Maybe B) -> IntMap A -> Property +prop_mapMaybeModel fun m = + mapMaybe (apply fun) m + === mapMaybeWithKey (const (apply fun)) m + +prop_mapMaybeWithKeyModel :: Fun (Key, A) (Maybe B) -> IntMap A -> Property +prop_mapMaybeWithKeyModel fun m = + mapMaybeWithKey (applyFun2 fun) m + === fromList (Maybe.mapMaybe fPair (toList m)) + where + fPair (k, v) = case apply fun (k, v) of + Nothing -> Nothing + Just v' -> Just (k, v') + +prop_mapEitherModel :: Fun A (Either B C) -> IntMap A -> Property +prop_mapEitherModel fun m = + mapEither (apply fun) m + === mapEitherWithKey (const (apply fun)) m + +prop_mapEitherWithKeyModel :: Fun (Key, A) (Either B C) -> IntMap A -> Property +prop_mapEitherWithKeyModel fun m = + mapEitherWithKey (applyFun2 fun) m + === ( mapMaybeWithKey (\k v -> maybeLeft (apply fun (k, v))) m + , mapMaybeWithKey (\k v -> maybeRight (apply fun (k, v))) m) + where + maybeLeft (Left v) = Just v + maybeLeft _ = Nothing + + maybeRight (Right v) = Just v + maybeRight _ = Nothing + prop_map :: Fun Int Int -> [(Int, Int)] -> Property -prop_map f ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_map f ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in map (apply f) m == fromList [ (a, apply f b) | (a,b) <- xs ] + in map (apply f) m === fromList [ (a, apply f b) | (a,b) <- xs ] prop_fmap :: Fun Int Int -> [(Int, Int)] -> Property -prop_fmap f ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_fmap f ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in fmap (apply f) m == fromList [ (a, apply f b) | (a,b) <- xs ] + in fmap (apply f) m === fromList [ (a, apply f b) | (a,b) <- xs ] prop_mapkeys :: Fun Int Int -> [(Int, Int)] -> Property -prop_mapkeys f ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_mapkeys f ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in mapKeys (apply f) m == (fromList $ List.nubBy ((==) `on` fst) $ reverse [ (apply f a, b) | (a,b) <- sort xs]) + in mapKeys (apply f) m === (fromList $ sortNubBy (compare `on` fst) $ reverse [ (apply f a, b) | (a,b) <- xs]) prop_splitModel :: Int -> [(Int, Int)] -> Property -prop_splitModel n ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_splitModel n ys = + let xs = sortNubBy (compare `on` fst) ys (l, r) = split n $ fromList xs - in valid l .&&. - valid r .&&. - toAscList l === sort [(k, v) | (k,v) <- xs, k < n] .&&. - toAscList r === sort [(k, v) | (k,v) <- xs, k > n] - -prop_splitRoot :: IMap -> Bool -prop_splitRoot s = loop ls && (s == unions ls) + in validProp l .&&. + validProp r .&&. + l === fromList (takeWhile ((< n) . fst) xs) .&&. + r === fromList (dropWhile ((<= n) . fst) xs) + +prop_splitLookupModel :: Int -> IMap -> Property +prop_splitLookupModel k m = + let (l, r) = split k m + in (l, lookup k m, r) === splitLookup k m + +prop_splitRoot :: IMap -> Property +prop_splitRoot s = loop ls .&&. (s === unions ls) where ls = splitRoot s loop [] = True @@ -1514,77 +1800,95 @@ prop_splitRoot s = loop ls && (s == unions ls) , y <- toList (unions rst) , x > y ] +-- 'isSubmapOf' and friends short-circuit all over the place, so test them a +-- lot more. +-- TODO: Is there a way to increase the number of tests that still allows +-- specifying a number at the command line? For example, multiply the +-- passed-in number of tests by 100? +increaseTests = withMaxSuccess 10000 + +prop_isSubmapOf :: IMap -> IMap -> Property +prop_isSubmapOf m1 m2 = increaseTests $ (m1 `isSubmapOf` m2) + === all (\(k, v) -> lookup k m2 == Just v) (toList m1) + +prop_isSubmapOfBy :: Fun (Int, Int) Bool -> IMap -> IMap -> Property +prop_isSubmapOfBy p m1 m2 = increaseTests $ isSubmapOfBy (curry (apply p)) m1 m2 + === all (\(k, v) -> member k m2 && apply p (v, m2 ! k)) (toList m1) + +prop_isProperSubmapOf :: IMap -> IMap -> Property +prop_isProperSubmapOf m1 m2 = increaseTests $ (m1 `isProperSubmapOf` m2) + === (size m1 < size m2 && m1 `isSubmapOf` m2) + +prop_isProperSubmapOfBy :: Fun (Int, Int) Bool -> IMap -> IMap -> Property +prop_isProperSubmapOfBy p m1 m2 = increaseTests $ isProperSubmapOfBy (curry (apply p)) m1 m2 + === (size m1 < size m2 && isSubmapOfBy (curry (apply p)) m1 m2) + prop_foldr :: Int -> [(Int, Int)] -> Property -prop_foldr n ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_foldr n ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in foldr (+) n m == List.foldr (+) n (List.map snd xs) && - foldr (:) [] m == List.map snd (List.sort xs) && - foldrWithKey (\_ a b -> a + b) n m == List.foldr (+) n (List.map snd xs) && - foldrWithKey (\k _ b -> k + b) n m == List.foldr (+) n (List.map fst xs) && - foldrWithKey (\k x xs -> (k,x):xs) [] m == List.sort xs + in foldr (+) n m === List.foldr (+) n (List.map snd xs) .&&. + foldr (:) [] m === List.map snd xs .&&. + foldrWithKey (\_ a b -> a + b) n m === List.foldr (+) n (List.map snd xs) .&&. + foldrWithKey (\k _ b -> k + b) n m === List.foldr (+) n (List.map fst xs) .&&. + foldrWithKey (\k x xs -> (k,x):xs) [] m === xs prop_foldr' :: Int -> [(Int, Int)] -> Property -prop_foldr' n ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_foldr' n ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in foldr' (+) n m == List.foldr (+) n (List.map snd xs) && - foldr' (:) [] m == List.map snd (List.sort xs) && - foldrWithKey' (\_ a b -> a + b) n m == List.foldr (+) n (List.map snd xs) && - foldrWithKey' (\k _ b -> k + b) n m == List.foldr (+) n (List.map fst xs) && - foldrWithKey' (\k x xs -> (k,x):xs) [] m == List.sort xs + in foldr' (+) n m === List.foldr (+) n (List.map snd xs) .&&. + foldr' (:) [] m === List.map snd xs .&&. + foldrWithKey' (\_ a b -> a + b) n m === List.foldr (+) n (List.map snd xs) .&&. + foldrWithKey' (\k _ b -> k + b) n m === List.foldr (+) n (List.map fst xs) .&&. + foldrWithKey' (\k x xs -> (k,x):xs) [] m === xs prop_foldl :: Int -> [(Int, Int)] -> Property -prop_foldl n ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_foldl n ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in foldl (+) n m == List.foldr (+) n (List.map snd xs) && - foldl (flip (:)) [] m == reverse (List.map snd (List.sort xs)) && - foldlWithKey (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) && - foldlWithKey (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) && - foldlWithKey (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs) + in foldl (+) n m === List.foldr (+) n (List.map snd xs) .&&. + foldl (flip (:)) [] m === reverse (List.map snd xs) .&&. + foldlWithKey (\b _ a -> a + b) n m === List.foldr (+) n (List.map snd xs) .&&. + foldlWithKey (\b k _ -> k + b) n m === List.foldr (+) n (List.map fst xs) .&&. + foldlWithKey (\xs k x -> (k,x):xs) [] m === reverse xs prop_foldl' :: Int -> [(Int, Int)] -> Property -prop_foldl' n ys = length ys > 0 ==> - let xs = List.nubBy ((==) `on` fst) ys +prop_foldl' n ys = + let xs = sortNubBy (compare `on` fst) ys m = fromList xs - in foldl' (+) n m == List.foldr (+) n (List.map snd xs) && - foldl' (flip (:)) [] m == reverse (List.map snd (List.sort xs)) && - foldlWithKey' (\b _ a -> a + b) n m == List.foldr (+) n (List.map snd xs) && - foldlWithKey' (\b k _ -> k + b) n m == List.foldr (+) n (List.map fst xs) && - foldlWithKey' (\xs k x -> (k,x):xs) [] m == reverse (List.sort xs) + in foldl' (+) n m === List.foldr (+) n (List.map snd xs) .&&. + foldl' (flip (:)) [] m === reverse (List.map snd xs) .&&. + foldlWithKey' (\b _ a -> a + b) n m === List.foldr (+) n (List.map snd xs) .&&. + foldlWithKey' (\b k _ -> k + b) n m === List.foldr (+) n (List.map fst xs) .&&. + foldlWithKey' (\xs k x -> (k,x):xs) [] m === reverse xs prop_foldrEqFoldMap :: IntMap Int -> Property prop_foldrEqFoldMap m = - foldr (:) [] m === Data.Foldable.foldMap (:[]) m + foldr (:) [] m === Foldable.foldMap (:[]) m prop_foldrWithKeyEqFoldMapWithKey :: IntMap Int -> Property prop_foldrWithKeyEqFoldMapWithKey m = foldrWithKey (\k v -> ((k,v):)) [] m === foldMapWithKey (\k v -> ([(k,v)])) m prop_FoldableTraversableCompat :: Fun A [B] -> IntMap A -> Property -prop_FoldableTraversableCompat fun m = foldMap f m === foldMapDefault f m +prop_FoldableTraversableCompat fun m = Foldable.foldMap f m === foldMapDefault f m where f = apply fun -prop_keysSet :: [(Int, Int)] -> Bool +#if MIN_VERSION_base(4,8,0) +prop_elem :: Int -> IMap -> Property +prop_elem v m = Foldable.elem v m === List.elem v (elems m) +#endif + +prop_keysSet :: [(Int, Int)] -> Property prop_keysSet xs = - keysSet (fromList xs) == IntSet.fromList (List.map fst xs) + keysSet (fromList xs) === IntSet.fromList (List.map fst xs) -prop_fromSet :: [(Int, Int)] -> Bool +prop_fromSet :: [(Int, Int)] -> Property prop_fromSet ys = - let xs = List.nubBy ((==) `on` fst) ys - in fromSet (\k -> fromJust $ List.lookup k xs) (IntSet.fromList $ List.map fst xs) == fromList xs - -newtype Identity a = Identity a - deriving (Eq, Show) - -instance Functor Identity where - fmap f (Identity a) = Identity (f a) - -instance Applicative Identity where - pure a = Identity a - Identity f <*> Identity a = Identity (f a) + let xs = sortNubBy (compare `on` fst) ys + in fromSet (\k -> fromJust $ List.lookup k xs) (IntSet.fromList $ List.map fst xs) === fromList xs prop_traverseWithKey_identity :: IntMap A -> Property prop_traverseWithKey_identity mp = mp === newMap @@ -1597,6 +1901,10 @@ prop_traverseWithKey_degrade_to_mapWithKey fun mp = g k v = Identity $ f k v Identity newMap = traverseWithKey g mp +-- While this isn't part of the IntMap API yet, it is still useful for testing traverseMaybeMissing +traverseMaybeWithKey :: Applicative f => (Key -> a -> f (Maybe b)) -> IntMap a -> f (IntMap b) +traverseMaybeWithKey = runWhenMissingAll . traverseMaybeMissing + prop_traverseMaybeWithKey_identity :: IntMap A -> Property prop_traverseMaybeWithKey_identity mp = mp === newMap where Identity newMap = traverseMaybeWithKey (\_ -> Identity . Just) mp @@ -1615,3 +1923,25 @@ prop_traverseMaybeWithKey_degrade_to_traverseWithKey fun mp = -- so this also checks the order of traversing is the same. where f k v = (show k, applyFun2 fun k v) g k v = fmap Just $ f k v + +prop_filterMissingEqFilterWithKey :: Fun (Int, A) Bool -> IntMap A -> Property +prop_filterMissingEqFilterWithKey fun m = + runIdentity (runWhenMissingAll (filterMissing f) m) === filterWithKey f m + where f = applyFun2 fun + +prop_filterAMissing_degrade_to_filterMissing :: Fun (Int, A) Bool -> IntMap A -> Property +prop_filterAMissing_degrade_to_filterMissing fun m = + runIdentity (runWhenMissingAll (filterAMissing (\k a -> Identity (f k a))) m) + === runIdentity (runWhenMissingAll (filterMissing f) m) + where f = applyFun2 fun + +prop_mapMissingEqMapWithKey :: Fun (Int, A) Int -> IntMap A -> Property +prop_mapMissingEqMapWithKey fun m = + runIdentity (runWhenMissingAll (mapMissing f) m) === mapWithKey f m + where f = applyFun2 fun + +prop_traverseMissing_degrade_to_mapMissing :: Fun (Int, A) Int -> IntMap A -> Property +prop_traverseMissing_degrade_to_mapMissing fun m = + runIdentity (runWhenMissingAll (traverseMissing (\k a -> Identity (f k a))) m) + === runIdentity (runWhenMissingAll (mapMissing f) m) + where f = applyFun2 fun diff --git a/containers/containers.cabal b/containers/containers.cabal index 42c01385e..0f7310bf7 100644 --- a/containers/containers.cabal +++ b/containers/containers.cabal @@ -43,9 +43,9 @@ Library Data.IntMap Data.IntMap.Lazy Data.IntMap.Strict - Data.IntMap.Strict.Internal Data.IntMap.Internal Data.IntMap.Internal.Debug + Data.IntMap.Merge.Internal Data.IntMap.Merge.Lazy Data.IntMap.Merge.Strict Data.IntSet.Internal @@ -77,6 +77,7 @@ Library if impl(ghc) other-modules: Utils.Containers.Internal.TypeError + Utils.Containers.Internal.IsList Data.Map.Internal.DeprecatedShowTree Data.IntMap.Internal.DeprecatedDebug diff --git a/containers/include/containers.h b/containers/include/containers.h index cd201ca3c..e622ea3d5 100644 --- a/containers/include/containers.h +++ b/containers/include/containers.h @@ -36,6 +36,8 @@ #ifdef __GLASGOW_HASKELL__ # define USE_ST_MONAD 1 # define USE_UNBOXED_ARRAYS 1 +# define USE_TYPE_FAMILIES 1 +# define USE_REWRITE_RULES 1 #endif #endif diff --git a/containers/src/Data/IntMap.hs b/containers/src/Data/IntMap.hs index 75855ded2..8e59e0f6e 100644 --- a/containers/src/Data/IntMap.hs +++ b/containers/src/Data/IntMap.hs @@ -14,8 +14,9 @@ ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 +-- Copyright : Documentation & Interface (c) Daan Leijen 2002 +-- Documentation (c) Andriy Palamarchuk 2008 +-- Documentation & Implementation (c) Jonathan S. 2016 -- License : BSD-style -- Maintainer : libraries@haskell.org -- Portability : portable @@ -36,20 +37,6 @@ -- > import Data.IntMap (IntMap) -- > import qualified Data.IntMap as IntMap -- --- The implementation is based on /big-endian patricia trees/. This data --- structure performs especially well on binary operations like 'union' --- and 'intersection'. However, my benchmarks show that it is also --- (much) faster on insertions and deletions when compared to a generic --- size-balanced map implementation (see "Data.Map"). --- --- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\", --- Workshop on ML, September 1998, pages 77-86, --- --- --- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve --- Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), --- October 1968, pages 514-534. --- -- Operation comments contain the operation time complexity in -- the Big-O notation . -- Many operations have a worst-case complexity of /O(min(n,W))/. diff --git a/containers/src/Data/IntMap/Internal.hs b/containers/src/Data/IntMap/Internal.hs index b8bc3ffb8..add17cfdc 100644 --- a/containers/src/Data/IntMap/Internal.hs +++ b/containers/src/Data/IntMap/Internal.hs @@ -1,27 +1,25 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} -#if __GLASGOW_HASKELL__ -{-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-} -{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE CPP, BangPatterns, EmptyDataDecls #-} +#if defined(__GLASGOW_HASKELL__) +{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving, MagicHash #-} +#endif + +#include "containers.h" + +#if USE_TYPE_FAMILIES +{-# LANGUAGE TypeFamilies #-} #endif #if !defined(TESTING) && defined(__GLASGOW_HASKELL__) {-# LANGUAGE Trustworthy #-} #endif -#if __GLASGOW_HASKELL__ >= 708 -{-# LANGUAGE TypeFamilies #-} -#endif {-# OPTIONS_HADDOCK not-home #-} -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Internal --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 --- (c) wren romano 2016 +-- Copyright : Documentation & Interface (c) Daan Leijen 2002 +-- Documentation (c) Andriy Palamarchuk 2008 +-- Documentation & Implementation (c) Jonathan "gereeter" S. 2020 -- License : BSD-style -- Maintainer : libraries@haskell.org -- Portability : portable @@ -43,17 +41,252 @@ -- This defines the data structures and core (hidden) manipulations -- on representations. -- +-- = Tree Structure +-- +-- This implementation uses a novel modification of /big-endian patricia trees/, structured +-- as a vantage-point tree under the XOR metric. +-- +-- = Derivation +-- +-- At its core, 'IntMap'\'s representation can be derived by a series of optimizations from +-- simpler structures: +-- +-- * A bitwise trie is compressed into a PATRICIA tree (a bitwise radix tree) by +-- merging series of nodes that have no branches. +-- * The prefix labels of PATRICIA tree nodes are represented implicitly by storing the minimum +-- and maximum keys in a subtree. +-- * Minima and maxima are only stored once, at the topmost location that they appear. +-- * Values are stored next to their associated keys, rather than in the leaves. +-- +-- Each of these steps is explained in detail below. +-- +-- == The basic integer map: the bitwise trie +-- +-- We are trying to create an efficient, simple mapping from integers to values. The most common +-- approaches are hash tables, which are not persistent (though we can come close with HAMTs), +-- and binary search trees, which work well, but don't use any special properties of the integer. +-- Thinking of integers not as numbers but as strings of bits, we use a /trie/, where a string is +-- interpreted as a series of instructions for which branch to take when navigating the tree. As +-- bits are particularly simple, so is the resulting structure: +-- +-- > data IntMap a = Bin (IntMap a) (IntMap a) | Tip a | Nil +-- +-- The `Bin` constructor represents a bitwise branch, and the `Tip` constructor comes after 64 +-- 64 `Bin` construtors in the tree (on a 64-bit machine). The associated basic operations navigate +-- the tree by reading a key bit by bit, at each node taking the branch associated with the current +-- bit: +-- +-- > lookup :: Int -> IntMap a -> Maybe a +-- > lookup k = go (finiteBitSize k - 1) +-- > where +-- > go b (Bin l r) = if testBit k b +-- > then go (b - 1) l +-- > else go (b - 1) r +-- > go _ (Tip x) = Just x +-- > go _ Nil = Nothing +-- > +-- > insert :: Int -> a -> IntMap a -> IntMap a +-- > insert k a = go (finiteBitSize k - 1) +-- > where +-- > go (-1) _ = Tip a +-- > go b (Bin l r) = if testBit k b +-- > then Bin (go (b - 1) l) r +-- > else Bin l (go (b - 1) r) +-- > go b _ = if testBit b k +-- > then Bin (go (b + 1) Nil) Nil +-- > else Bin Nil (go (b + 1) Nil) +-- +-- 'delete' follows similarly, and the uniform structure means that even 'union' isn't too hard, +-- a welcome fact given the complexity of merging binary search trees. Unfortunately, this +-- approach is extremely slow and space-inefficient. To see why, look at the tree structure +-- for @'singleton' 5 "hello"@: +-- +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─0─┐ +-- > └─1─┐ +-- > └─0─┐ +-- > └─1─ "hello" +-- +-- Note that, for brevity, the word size in this diagram is 16 bits. It would be 4 times longer +-- for a 64-bit system. In this atrocious tree structure, there is one pointer for every bit, a +-- 64-fold explosion in space. Arguably worse is the fact that every single 'lookup', +-- 'Data.IntMap.Lazy.insert', or 'delete' must traverse 64 pointers, resulting in 64 cache misses +-- and a corresponding slowdown. +-- +-- == Path compression: PATRICIA trees and the previous version of 'Data.IntMap' +-- +-- To reduce space usage, we compress nodes that only have one child. Since they form a +-- linear chain, we can concatenate the bits within that chain, recording which branches would be +-- taken. For example, again temporarily shortening the word size to 16 bits: +-- +-- >>> singleton 5 "hello" +-- └─0000000000000101─ "hello" +-- +-- >>> fromList [(1, "1"), (4, "4"), (5, "5")] +-- └─0000000000000─┐ +-- ├─001─ "1" +-- └─10─┐ +-- ├─0─ "4" +-- └─1─ "5" +-- +-- This is much more space-efficient, and the basic operations, while more complicated, are still +-- straightforward. In Haskell, the structure is +-- +-- > type Prefix = Int +-- > type Mask = Int +-- > data IntMap a = Bin Prefix Mask (IntMap a) (IntMap a) | Tip Int a | Nil +-- +-- The @Mask@ tells how long the @Prefix@ is, and the @Int@ in the @Tip@ nodes encodes the +-- remaining bits. This representation, known as the big-endian PATRICIA tree, is what the +-- previous iteration of 'IntMap' used. +-- +-- == Implicit prefixes: a simpler representation +-- +-- In the PATRICIA tree representation above, we explicitly stored the common prefix of all the +-- keys in a subtree. However, this prefix is not needed if we know the minimum and maximum keys. +-- The common prefix of a set of keys is the same as the common prefix of the minimum and maximum. +-- Replacing the @Prefix@, @Mask@ pair with a minimum and maximum, we get +-- +-- > type MinBound = Int +-- > type MaxBound = Int +-- > data IntMap a = Bin MinBound MaxBound (IntMap a) (IntMap a) | Tip Int a | Nil +-- +-- The tree structure looks identical, just with different labels on the edges: +-- +-- >>> singleton 5 "hello" +-- └─5─ "hello" +-- +-- >>> fromList [(1, "1"), (4, "4"), (5, "5")] +-- └─(1,5)─┐ +-- ├─1─ "1" +-- └─(4,5)─┐ +-- ├─4─ "4" +-- └─5─ "5" +-- +-- Traversing this tree efficiently is a bit more difficult, but still possible. See 'xor' for +-- details. Moreover, since the tree contains exact minima and maxima, 'lookup' can already be +-- more efficient than in a PATRICIA tree. Even if a key matches the prefix of common bits, if the +-- key is out of the bounds of a subtree, a search can terminate early with 'Nothing'. However, +-- there are bigger gains to be had. +-- +-- == Removing redundancy +-- +-- The above representation stores many keys repeatedly. In the @{1,4,5}@ example, 1 was stored +-- twice, 4 was stored twice, and 5 was stored three times. The minimum and maximum keys of a +-- tree are necessarily keys stored in that tree and moreover are minima and maxima of subtrees. +-- In the @{1,4,5}@ example, we know from the root node that the minimum is 1 and the maximum is +-- 5. At the first branch, we split the set into two parts, @{1}@ and @{4,5}@. However, the +-- minimum of the set of lesser keys is equal to the minimum of the original set. Similarly, the +-- maximum of the set of greater keys is equal to the maximum of the original set. +-- +-- We can restructure the tree to store only one new value at each branch, removing the redundancy. +-- In nodes storing a set of lesser keys, we already know the minimum when traversing the tree +-- downward, so we only need to store the new maximum. In nodes storing a set of greater keys, we +-- know the maximum and store the new minimum. The root still needs both the minimum and the +-- maximum, so we need an extra layer to store that information: +-- +-- > type Bound = Int +-- > data IntMap a = Empty | NonEmpty Bound (Node a) +-- > data Node a = Bin Bound (Node a) (Node a) | Tip a +-- +-- The trees are no longer quite as easy to read at a glance, since keys are no longer visible in +-- order at the leaves. It can be difficult to tell the difference at a glance between a node +-- storing a minimum and a node storing a maximum. (The actual implementation uses phantom types +-- to ensure no code gets this wrong.) +-- +-- >>> singleton 5 "hello" +-- 5 +-- └─ "hello" +-- +-- >>> fromList [(1, "1"), (4, "4"), (5, "5")] +-- 1 +-- └─5─┐ +-- ├─ "1" +-- └─4─┐ +-- ├─ "4" +-- └─ "5" +-- +-- It may be easier to visualize what is happening here if we draw our trees +-- with maxima coming after their children: +-- +-- >>> fromList [(1, "1"), (4, "4"), (5, "5")] +-- 1 +-- │ ┌─ "1" +-- │ ├─4─┐ +-- │ │ ├─ "4" +-- │ │ └─ "5" +-- └─5─┘ +-- +-- Although the nonuniform tree structure results in more complex code, we save a word in each +-- node. +-- +-- == Moving the values upward +-- +-- The above change removed the redundancy in keys perfectly, so each key is stored only once. +-- However, the values are still stored in leaves, now far away from their associated keys. +-- There is no reason this has to be true now that each keys has a unique location in the tree. We +-- simplify by moving the values next to their keys: +-- +-- > data IntMap a = Empty | NonEmpty Bound a (Node a) +-- > data Node a = Bin Bound a (Node a) (Node a) | Tip +-- +-- Although nodes still switch between minima and maxima, they can be visualized and manipulated +-- more cleanly since it is clear which keys are tied to which values. +-- +-- >>> singleton 5 "hello" +-- 5 "hello" +-- └╼ +-- +-- >>> fromList [(1, "1"), (4, "4"), (5, "5")] +-- 1 "1" +-- │ ┌╼ +-- │ ├─4─┐ "4" +-- │ │ ├╼ +-- │ │ └╼ +-- └─5─┘ "5" +-- +-- This simpler representation translates to even more savings in both space and time. Since the +-- leaves no longer store any information, GHC will create a single static @Tip@ object and reuse +-- it for all leaves, the equivalent of representing leaves with a null pointer. This saves on +-- allocations and the metadata necessary for garbage collection and lazy evaluation. +-- Additionally, successful lookups can terminate as soon as they see the correct key instead of +-- dereferencing a chain of pointers all the way to the leaves. This means fewer cache misses and +-- a shorter loop. +-- +-- = References and Further Reading +-- +-- Morrison introduced PATRICIA trees in: +-- +-- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve Information Coded In Alphanumeric/\" +-- Journal of the ACM, 15(4), October 1968, pages 514-534. +-- +-- Okasaki and Gill proposed using them in a functional context and provided implementations, +-- benchmarks, and discussion in: +-- +-- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\", +-- Workshop on ML, September 1998, pages 77-86, +-- . +-- +-- Kmett proposed replacing explicit prefixes with min/max pairs in: +-- +-- * Edward Kmett, \"/Revisiting Matrix Multiplication, Part IV: IntMap!?/\", +-- School of Haskell, 25 August 2013, +-- . +-- -- @since 0.5.9 ----------------------------------------------------------------------------- --- [Note: INLINE bit fiddling] --- ~~~~~~~~~~~~~~~~~~~~~~~~~~~ --- It is essential that the bit fiddling functions like mask, zero, branchMask --- etc are inlined. If they do not, the memory allocation skyrockets. The GHC --- usually gets it right, but it is disastrous if it does not. Therefore we --- explicitly mark these functions INLINE. - - -- [Note: Local 'go' functions and capturing] -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -- Care must be taken when using 'go' function which captures an argument. @@ -62,125 +295,93 @@ -- must be checked for increased allocation when creating and modifying such -- functions. +module Data.IntMap.Internal ( + -- * Map Types + IntMap(..) + , L, R + , IntMap_(..) + , Node(..) + , NonEmptyIntMap_(..) + + -- ** Key Manipulation + , Key + , UKey + , box + , unbox + , Bound(..) + , boundUKey + , BoundOrdering(..) + , xor + , xorBounds + , compareMinBound + , compareMaxBound + , ltMSB + , compareMSB + , boundsDisjoint + , minToMax + , maxToMin --- [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%. + -- * Construction + , empty -module Data.IntMap.Internal ( - -- * Map type - IntMap(..), Key -- instance Eq,Show + -- ** From Unordered Lists + , fromListLazy - -- * Operators - , (!), (!?), (\\) + -- ** From Ascending Lists + , BuildStack(..) + , pushBuildStack + , completeBuildStack + + -- * Insertion + , insertLazy + , insertWithEval + , insertMinL + , insertMaxR + + -- * Deletion\/Update + , delete + , deleteMinL + , deleteMaxR -- * Query - , null - , size - , member - , notMember + -- ** Lookup , lookup + , (!?) + , (!) , findWithDefault + , member + , notMember , lookupLT , lookupGT , lookupLE , lookupGE - , disjoint - - -- * Construction - , empty - , singleton - - -- ** Insertion - , insert - , insertWith - , insertWithKey - , insertLookupWithKey - -- ** Delete\/Update - , delete - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - , alterF + -- ** Size + , null + , size -- * Combine - -- ** Union , union - , unionWith - , unionWithKey , unions - , unionsWith + , unionDisjointL + , unionDisjointR -- ** Difference , difference - , differenceWith - , differenceWithKey + , (\\) -- ** Intersection , intersection - , intersectionWith - , intersectionWithKey - - -- ** General combining function - , SimpleWhenMissing - , SimpleWhenMatched - , runWhenMatched - , runWhenMissing - , merge - -- *** @WhenMatched@ tactics - , zipWithMaybeMatched - , zipWithMatched - -- *** @WhenMissing@ tactics - , mapMaybeMissing - , dropMissing - , preserveMissing - , mapMissing - , filterMissing - - -- ** Applicative general combining function - , WhenMissing (..) - , WhenMatched (..) - , mergeA - -- *** @WhenMatched@ tactics - -- | The tactics described for 'merge' work for - -- 'mergeA' as well. Furthermore, the following - -- are available. - , zipWithMaybeAMatched - , zipWithAMatched - -- *** @WhenMissing@ tactics - -- | The tactics described for 'merge' work for - -- 'mergeA' as well. Furthermore, the following - -- are available. - , traverseMaybeMissing - , traverseMissing - , filterAMissing - - -- ** Deprecated general combining function - , mergeWithKey - , mergeWithKey' + + -- ** Disjoint + , disjoint -- * Traversal -- ** Map - , map - , mapWithKey - , traverseWithKey - , traverseMaybeWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeys - , mapKeysWith - , mapKeysMonotonic + , mapLazy + , mapStrict_ + , mapNodeStrict -- * Folds , foldr @@ -200,21 +401,25 @@ module Data.IntMap.Internal ( , keys , assocs , keysSet - , fromSet -- ** Lists , toList - , fromList - , fromListWith - , fromListWithKey - -- ** Ordered lists + -- ** Ordered Lists , toAscList , toDescList - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList + + -- ** Internal Manipulation + , binL + , binR + , binNodeMapL + , binMapNodeR + , extractBinL + , extractBinR + , l2rMap + , r2lMap + , nodeToMapL + , nodeToMapR -- * Filter , filter @@ -223,19 +428,15 @@ module Data.IntMap.Internal ( , withoutKeys , partition , partitionWithKey - - , mapMaybe - , mapMaybeWithKey - , mapEither - , mapEitherWithKey - , split , splitLookup , splitRoot -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy + , isSubmapOf + , isSubmapOfBy + , isProperSubmapOf + , isProperSubmapOfBy -- * Min\/Max , lookupMin @@ -246,288 +447,486 @@ module Data.IntMap.Internal ( , deleteMax , deleteFindMin , deleteFindMax - , updateMin - , updateMax - , updateMinWithKey - , updateMaxWithKey , minView , maxView , minViewWithKey , maxViewWithKey +) where - -- * Debugging - , showTree - , showTreeWith - - -- * Internal types - , Mask, Prefix, Nat - - -- * Utility - , natFromInt - , intFromNat - , link - , linkWithMask - , bin - , binCheckLeft - , binCheckRight - , zero - , nomatch - , match - , mask - , maskW - , shorter - , branchMask - , highestBitMask - - -- * Used by "IntMap.Merge.Lazy" and "IntMap.Merge.Strict" - , mapWhenMissing - , mapWhenMatched - , lmapWhenMissing - , contramapFirstWhenMatched - , contramapSecondWhenMatched - , mapGentlyWhenMissing - , mapGentlyWhenMatched - ) where +import Control.DeepSeq (NFData(..)) -#if MIN_VERSION_base(4,8,0) -import Data.Functor.Identity (Identity (..)) -import Control.Applicative (liftA2) -#else -import Control.Applicative (Applicative(pure, (<*>)), (<$>), liftA2) -import Data.Monoid (Monoid(..)) -import Data.Traversable (Traversable(traverse)) -import Data.Word (Word) -#endif +import Data.Maybe (fromMaybe) +import qualified Data.List (foldl') +import qualified Data.Foldable (Foldable(..)) #if MIN_VERSION_base(4,9,0) -import Data.Semigroup (Semigroup(stimes)) -#endif -#if !(MIN_VERSION_base(4,11,0)) && MIN_VERSION_base(4,9,0) -import Data.Semigroup (Semigroup((<>))) +#if MIN_VERSION_base(4,11,0) +import Data.Semigroup (stimes) +#else +import Data.Semigroup (Semigroup(..)) #endif -#if MIN_VERSION_base(4,9,0) import Data.Semigroup (stimesIdempotentMonoid) import Data.Functor.Classes #endif -import Control.DeepSeq (NFData(rnf)) -import Data.Bits -import qualified Data.Foldable as Foldable +#if defined(__GLASGOW_HASKELL__) #if !MIN_VERSION_base(4,8,0) -import Data.Foldable (Foldable()) +import Data.Functor ((<$)) #endif -import Data.Maybe (fromMaybe) import Data.Typeable -import Prelude hiding (lookup, map, filter, foldr, foldl, null) - -import Data.IntSet.Internal (Key) -import qualified Data.IntSet.Internal as IntSet -import Utils.Containers.Internal.BitUtil -import Utils.Containers.Internal.StrictPair - -#if __GLASGOW_HASKELL__ import Data.Data (Data(..), Constr, mkConstr, constrIndex, Fixity(Prefix), DataType, mkDataType) -import GHC.Exts (build) -#if !MIN_VERSION_base(4,8,0) -import Data.Functor ((<$)) +import Text.Read +import GHC.Exts (Int(..), Int#) #endif #if __GLASGOW_HASKELL__ >= 708 -import qualified GHC.Exts as GHCExts +import qualified Utils.Containers.Internal.IsList as IsList #endif -import Text.Read -#endif -import qualified Control.Category as Category -#if __GLASGOW_HASKELL__ >= 709 -import Data.Coerce +#if USE_REWRITE_RULES +import GHC.Exts (build) #endif +#if !MIN_VERSION_base(4,8,0) +import Data.Word (Word) +import Data.Monoid (Monoid(..)) +import Data.Traversable (Traversable(..)) +import Control.Applicative (Applicative(..), (<$>)) +#endif +import Control.Applicative (liftA2, liftA3) + +import qualified Data.Bits (xor) + +import qualified Data.IntSet (IntSet, fromDistinctAscList, member, notMember) +import Utils.Containers.Internal.StrictPair (StrictPair(..)) + +import Prelude hiding (foldr, foldl, filter, lookup, null, map, min, max) + +-- These definitions are the only things defining that this map applies to +-- 'Int' keys. Using 'Int8' or 'Word' or some other two's complement integer +-- type would produce a working, equally efficient map from that type. +type Key = Int +#if defined(__GLASGOW_HASKELL__) +-- | An unboxed, unlifted version of 'Key', used for manual worker-wrapper +-- transformations. @WithKey@ functions, particularly those that call their +-- combination function on many keys, should be @INLINE@ and written in terms +-- of a @WithUKey@ variant where the combination function takes a 'UKey'. +type UKey = Int# + +-- | Convert a 'UKey' into the equivalent 'Key'. +{-# INLINE box #-} +box :: UKey -> Key +box = I# + +-- | Convert a 'Key' into the equivalent 'UKey'. +{-# INLINE unbox #-} +unbox :: Key -> UKey +unbox (I# x) = x +#else +-- | Under GHC, this would be used to force a worker-wrapper transformation on +-- @WithKey@ functions. However, since this compiler doesn't support unlifted +-- types, there is no way to do this manually. +type UKey = Key + +-- | Convert a 'UKey' into the equivalent 'Key'. +{-# INLINE box #-} +box :: UKey -> Key +box = id + +-- | Convert a 'Key' into the equivalent 'UKey'. +{-# INLINE unbox #-} +unbox :: Key -> UKey +unbox = id +#endif --- A "Nat" is a natural machine word (an unsigned Int) -type Nat = Word +{-# INLINE i2w #-} +i2w :: Key -> Word +i2w = fromIntegral + +-- | XOR a key with a bound for the purposes of navigation within the tree. +-- +-- The navigation process is as follows. Suppose we are looking up a key @k@ in a tree. We know +-- the minimum key in the tree, @min@, and the maximum key, @max@. Represented in binary: +-- +-- > shared prefix bit to split on +-- > /----------\ / +-- > min: 010010010101 0 ???????? +-- > max: 010010010101 1 ???????? +-- > k: 010010010101 ? ???????? +-- +-- To figure out in which subtree might contain @k@, we need to know whether the bit to split on +-- is zero or one. Now, if it is zero, then +-- +-- > xor k min: 000000000000 0 ???????? +-- > xor k max: 000000000000 1 ???????? +-- +-- If it is one: +-- +-- > xor k min: 000000000000 1 ???????? +-- > xor k max: 000000000000 0 ???????? +-- +-- Therefore, the splitting bit is set iff @'xor' k min > 'xor' k max@. Put another way, the key +-- shares more bits with the bound that it is closer to under the XOR metric, since exclusive or +-- maps shared bits to zero. The metric perspective also makes it clear why this works unmodified +-- in the presence of negative numbers, despite storing negative numbers (with a set sign bit) in +-- the left branch normally identified with an unset bit. As long as the comparison is done +-- unsigned (metrics are always nonnegative), negative integers will be closer to other negative +-- integers than they are to nonnegative integers. +{-# INLINE xor #-} +xor :: Key -> Bound t -> Word +xor a (Bound b) = Data.Bits.xor (i2w a) (i2w b) + +-- | XOR the minimum and maximum keys of a tree. The most significant bit of the result indicates +-- which bit to split on for navigation and is useful in merging maps to tell whether nodes from +-- different maps branch at the same time. +{-# INLINE xorBounds #-} +xorBounds :: Bound L -> Bound R -> Word +xorBounds (Bound min) (Bound max) = Data.Bits.xor (i2w min) (i2w max) + +-- | Check if two bounds necessarily bound non-overlapping ranges of keys (i.e. +-- the minimum is greater than the maximum). This is used in merges, indicating +-- when the subtrees from the two maps become unrelated, never sharing another +-- key. See 'Data.IntMap.Merge.Internal' for further details on its use. +{-# INLINE boundsDisjoint #-} +boundsDisjoint :: Bound L -> Bound R -> Bool +boundsDisjoint (Bound min) (Bound max) = min > max + +-- These are uninhabited to ensure that they are only used as type parameters. +-- | A type tag denoting 'Node's found on left branches or minimum 'Bound's. +data L +-- | A type tag denoting 'Node's found on right branches or maximum 'Bound's. +data R + +#if USE_TYPE_FAMILIES +-- TODO: If we are relying on GHC features anyway, L and R could be a new kind. + +-- | A 'Key' stored somewhere within an 'IntMap' with a tag representing its role within that map. +-- Following the left-to-right ascending order of keys in the map, a @'Bound' 'L'@ serves the +-- role of a minimum (and so must be less than all other keys in the subtree where it was found), +-- and a @'Bound' 'R'@ serves the role of maximum (and so must be greater than all other keys in +-- the subtree where it was found). +-- +-- Because the tag represents the relationship between a 'Bound' and other keys in a map, the tag +-- will typically only change or be stripped away in a base case where the 'Bound' is known to be +-- in a trivial singleton subtree with only that one key. Therefore, 'boundKey', 'minToMax', and +-- 'maxToMin', which strip or change the tag, should be rare. +newtype Bound t = Bound { boundKey :: Key } deriving (Eq, Ord, Show) + +-- | The opposite direction of a tag. This is necessary as the 'Bound's in 'IntMap_' and 'Node' +-- are opposite each other, since they form the initial min/max bracket on the map as a whole. +-- This is used in 'Node', where a left branch inherits the minimum from its parent, so needs +-- to specify a new maximum (a right bound), and symmetrically in a right branch. +type family Flipped t +type instance Flipped L = R +type instance Flipped R = L +#else +-- | A 'Key' stored somewhere within an 'IntMap'. Since this was compiled without type families, +-- the associated tag can't be tracked properly and is therefore thrown away. +type Bound t = Bound_ +newtype Bound_ = Bound { boundKey :: Key } deriving (Eq, Ord, Show) +-- This, like L and R, is uninhabited to ensure that it is only used as a type parameter +data Flipped t +#endif -natFromInt :: Key -> Nat -natFromInt = fromIntegral -{-# INLINE natFromInt #-} +-- | Extract the key from a 'Bound' and unbox it. Identical in functionality to +-- 'boundKey'. +{-# INLINE boundUKey #-} +boundUKey :: Bound t -> UKey +boundUKey b = unbox (boundKey b) -intFromNat :: Nat -> Key -intFromNat = fromIntegral -{-# INLINE intFromNat #-} +data BoundOrdering = InBound | OutOfBound | Matched deriving (Eq) -{-------------------------------------------------------------------- - Types ---------------------------------------------------------------------} +{-# INLINE compareMinBound #-} +compareMinBound :: Key -> Bound L -> BoundOrdering +compareMinBound k (Bound min) + | k > min = InBound + | k < min = OutOfBound + | otherwise = Matched +{-# INLINE compareMaxBound #-} +compareMaxBound :: Key -> Bound R -> BoundOrdering +compareMaxBound k (Bound max) + | k < max = InBound + | k > max = OutOfBound + | otherwise = Matched -- | A map of integers to values @a@. +newtype IntMap a = IntMap (IntMap_ L a) deriving (Eq) + +-- | A self-contained tree mapping integers to values @a@. This is tagged according to the type of +-- its root 'Node'. Although the 'L' form is the only one used at the top level, both types are +-- used as intermediates in cases where 'Node' isn't applicable. This happens primarily when the +-- bound that forms the external context for a 'Node' has been deleted. See also 'binL' and +-- 'binR', which allow 'Bin'-like combination of 'IntMap_'s. +-- +-- Unlike a 'Node', which inherits one of its bounds from its parent, an 'IntMap_' stores both of +-- its bounds, one in the 'NonEmpty' constructor and one in the top-level 'Node'. This allows it +-- to be meaningfully navigated without additional external context. +-- +-- Invariants: +-- +-- * The keys within a tree are in order. Specifically, all descendants of wherever a minimum +-- (@'Bound' 'L'@) is stored must contain keys strictly greater than that minimum, and all +-- descendants of wherever a maximum (@'Bound' 'R'@) is stored must contain keys stricly +-- less than that maximum. +-- * All keys within the left branch of a 'Bin' 'Node' must be stricly less than all keys in the +-- right branch of that same 'Node'. Notably, this is true regardless of the signs of the keys; +-- in a map containing negative and nonnegative keys, the top-level 'Bin' will split with the +-- negative keys on the left and the nonnegative keys on the right. +-- * All keys within the left branch of a 'Bin' 'Node' must share a longer bit prefix with the +-- minimum bound for the 'Node' than with the maximum bound for the 'Node', and vice versa for +-- the right branch: all keys within the right branch of a 'Bin' 'Node' must share a longer bit +-- prefix with the maximum bound for the 'Node' than with the minimum bound for the 'Node'. +-- Put another way, under the XOR metric, keys in the left branch are closer to the minimum +-- bound than the maximum bound (@'xor' k min < 'xor' k max@), while keys in the right branch +-- are closer to the maximum bound than the minimum bound. (@'xor' k min > 'xor' k max@). +-- +-- The ordering requirements on keys mean that any key will share whatever prefix bits the +-- bounds also share. This invariant implies that branches are taken based on the immediately +-- following bit. When the minimum and maximum are both of the same sign, the 0 branch is lesser +-- and therefore assigned to the left side. For example, when splitting between 4 (@0100@) and +-- 7 (@0111@), which disagree in the 2's place, 5 (@0101@) would be found on the left side since +-- it has a 0 in the 2's place; similarly, -7 (@1001@) and -4 (@1100@) disagree in the 4's +-- place, and -5 (@1011@) would be found on the left side since it has a 0 in the 4's place. +-- However, when the minimum and maximum disagree in sign, they differ in the first bit, the +-- sign bit. A 1 in the sign bit indicates a negative number, so in this case the 1 branch is +-- lesser and therefore assigned to the left side. For example, when splitting between -7 +-- (@1001@) and 7 (@0111@), which disagree in the sign bit, 5 (@0101@) would be found on the +-- /right/ side since it has a 0 in the sign bit. +-- +-- Note that one of these minimum and maximum bounds is inherited from the parent of the 'Node'. +-- The terms do not refer to the minimum and maximum keys actually stored somewhere within the +-- 'Node' or in one of its descendants. +-- +-- These invariants imply a unique tree structure for a given set of keys. In fact, they are +-- overspecified: the second invariant follows from the other two. To check these invariants, +-- use 'Data.IntMap.Internal.Debug.valid'. +data IntMap_ t a = NonEmpty {-# UNPACK #-} !(Bound t) a !(Node t a) | Empty deriving (Eq) + +-- | A node within a tree mapping integers to value @a@. Unlike an 'IntMap_', a 'Node' cannot +-- stand fully on its own; it must have some context defining where the splits between left and +-- right are. A @'Node' 'L'@ is typically found on the left branch of another 'Node', so inherits +-- its minimum bound from its parent 'Node', storing only a new maximum. Similarly, a @'Node' 'R'@ +-- is typically found on the right branch of another 'Node', so inherits its maximum bound from +-- its parent 'Node', storing only a new minimum. See 'IntMap_' for further discussion of the tree +-- structure. +-- +-- Because of its incompleteness, functions that navigate 'Node's will typically pass them in two +-- arguments, one providing the missing bound and one providing the actual 'Node'. This isn't +-- universally true. Some functions, like 'map' or 'filter', preserve the overall branch structure +-- and don't need to understand the criteria used to choose between left and right. Others, like +-- 'lookup', only need to understand the branching criteria with regards to a single, fixed key, +-- and can instead pass the XOR-distance between that key and the missing bound. +data Node t a = Bin {-# UNPACK #-} !(Bound (Flipped t)) a !(Node L a) !(Node R a) | Tip deriving (Eq, Show) + +-- | The non-empty case of 'IntMap_'. +-- +-- Although this may be exposed and used for its own sake (TODO: make 'IntMap_' +-- contain an @UNPACK@ed 'NonEmptyIntMap_' field instead of copy-pasting +-- fields between the two structures), its primary function currently is as an +-- intermediate in routines that involve deletion. +-- +-- A 'Node' is largely meaningless without the context of one external bound +-- (minima for 'L' nodes and maxima for 'R' nodes). When that external bound +-- gets deleted, the 'Node' needs to be rearranged, with a new bouond pulled +-- up and out to the top level. This transformation is implemented in +-- 'deleteMinL' for when the minimum bound is gone and in 'deleteMaxR' when +-- the minimum bound is gone. Although a 'Node' could be 'Tip' and so would +-- most naturally be turned into a possibly-empty 'IntMap_' (and this is what +-- 'nodeToMapL' and 'nodeToMapR' do), it is important to use a product type for +-- the intermediates in the recursive case so that GHC's constructed product +-- result analysis (CPR) can unpack the 'NonEmptyIntMap_' into multiple +-- returned values on the stack. +-- +-- See 'IntMap_' for further discussion of the tag type parameter and the +-- invariants that must hold for this structure. +data NonEmptyIntMap_ t a = NE {-# UNPACK #-} !(Bound t) a !(Node t a) --- See Note: Order of constructors -data IntMap a = Bin {-# UNPACK #-} !Prefix - {-# UNPACK #-} !Mask - !(IntMap a) - !(IntMap a) --- Fields: --- prefix: The most significant bits shared by all keys in this Bin. --- mask: The switching bit to determine if a key should follow the left --- or right subtree of a 'Bin'. --- Invariant: Nil is never found as a child of Bin. --- Invariant: The Mask is a power of 2. It is the largest bit position at which --- two keys of the map differ. --- Invariant: Prefix is the common high-order bits that all elements share to --- the left of the Mask bit. --- Invariant: In (Bin prefix mask left right), left consists of the elements that --- don't have the mask bit set; right is all the elements that do. - | Tip {-# UNPACK #-} !Key a - | Nil - -type Prefix = Int -type Mask = Int - - --- Some stuff from "Data.IntSet.Internal", for 'restrictKeys' and --- 'withoutKeys' to use. -type IntSetPrefix = Int -type IntSetBitMap = Word - -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' +#if MIN_VERSION_base(4,9,0) +-- | @since 0.5.9 +instance Eq1 IntMap where + liftEq eq (IntMap m1) (IntMap m2) = liftEq eq m1 m2 + +instance Eq1 (IntMap_ t) where + liftEq eq (NonEmpty min1 minV1 root1) (NonEmpty min2 minV2 root2) + = min1 == min2 && eq minV1 minV2 && liftEq eq root1 root2 + liftEq _ Empty Empty = True + liftEq _ _ _ = False + +instance Eq1 (Node t) where + liftEq eq (Bin k1 v1 l1 r1) (Bin k2 v2 l2 r2) + = k1 == k2 && eq v1 v2 && liftEq eq l1 l2 && liftEq eq r1 r2 + liftEq _ Tip Tip = True + liftEq _ _ _ = False +#endif -(!) :: IntMap a -> Key -> a -(!) m k = find k m +instance Ord a => Ord (IntMap a) where + compare m1 m2 = compare (toList m1) (toList m2) + m1 <= m2 = toList m1 <= toList m2 --- | /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 +#if MIN_VERSION_base(4,9,0) +-- | @since 0.5.9 +instance Ord1 IntMap where + liftCompare cmp m1 m2 = liftCompare (liftCompare cmp) (toList m1) (toList m2) +#endif -(!?) :: IntMap a -> Key -> Maybe a -(!?) m k = lookup k m --- | Same as 'difference'. -(\\) :: IntMap a -> IntMap b -> IntMap a -m1 \\ m2 = difference m1 m2 +instance Show a => Show (IntMap a) where + showsPrec precedence m = showParen (precedence > 10) (showString "fromList " . shows (toList m)) -infixl 9 !?,\\{-This comment teaches CPP correct behaviour -} +#if MIN_VERSION_base(4,9,0) +-- | @since 0.5.9 +instance Show1 IntMap where + liftShowsPrec innerShowsPrec innerShowList precedence m + = showsUnaryWith listShowsPrec "fromList" precedence (toList m) + where + listShowsPrec = liftShowsPrec pairShowsPrec pairShowList + pairShowsPrec = liftShowsPrec innerShowsPrec innerShowList + pairShowList = liftShowList innerShowsPrec innerShowList +#endif -{-------------------------------------------------------------------- - Types ---------------------------------------------------------------------} +instance Read a => Read (IntMap a) where +#if defined(__GLASGOW_HASKELL__) + -- ReadPrec is more efficient than ReadS, so use it if possible + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + fromListLazy <$> readPrec + readListPrec = readListPrecDefault +#else + readsPrec precedence = readParen (precedence > 10) $ \str -> do + ("fromList", str') <- lex str + first fromList <$> reads str' +#endif -instance Monoid (IntMap a) where - mempty = empty - mconcat = unions -#if !(MIN_VERSION_base(4,9,0)) - mappend = union +#if MIN_VERSION_base(4,9,0) +-- | @since 0.5.9 +instance Read1 IntMap where +#if defined(__GLASGOW_HASKELL__) && MIN_VERSION_base(4,10,0) + liftReadPrec innerOne innerList = readData (readUnaryWith listOne "fromList" fromListLazy) + where + listOne = liftReadPrec pairOne pairList + pairOne = liftReadPrec innerOne innerList + pairList = liftReadListPrec innerOne innerList + liftReadListPrec = liftReadListPrecDefault #else - mappend = (<>) + liftReadsPrec innerOne innerList = readsData (readsUnaryWith listOne "fromList" fromListLazy) + where + listOne = liftReadsPrec pairOne pairList + pairOne = liftReadsPrec innerOne innerList + pairList = liftReadList innerOne innerList +#endif +#endif --- | @since 0.5.7 -instance Semigroup (IntMap a) where - (<>) = union - stimes = stimesIdempotentMonoid +instance Functor IntMap where + fmap = mapLazy + +#if defined(__GLASGOW_HASKELL__) + a <$ (IntMap m) = IntMap (a <$ m) +#endif + +instance Functor (IntMap_ t) where + fmap = mapLazy_ + +#if defined(__GLASGOW_HASKELL__) + _ <$ Empty = Empty + a <$ NonEmpty min _ node = NonEmpty min a (a <$ node) +#endif + +instance Functor (Node t) where + fmap = mapNodeLazy + +#if defined(__GLASGOW_HASKELL__) + _ <$ Tip = Tip + a <$ Bin k _ l r = Bin k a (a <$ l) (a <$ r) #endif -- | Folds in order of increasing key. -instance Foldable.Foldable IntMap where - fold = go - where go Nil = mempty - go (Tip _ v) = v - go (Bin _ m l r) - | m < 0 = 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 _ m l r) - | m < 0 = go r `mappend` go l - | otherwise = go l `mappend` go r - {-# INLINE foldMap #-} - foldl' = foldl' - {-# INLINE foldl' #-} - foldr' = foldr' - {-# INLINE foldr' #-} +instance Data.Foldable.Foldable IntMap where + {-# INLINE foldMap #-} + foldMap f = start + where + start (IntMap Empty) = mempty + start (IntMap (NonEmpty _ minV root)) = f minV `mappend` goL root + + goL Tip = mempty + goL (Bin _ maxV l r) = goL l `mappend` goR r `mappend` f maxV + + goR Tip = mempty + goR (Bin _ minV l r) = f minV `mappend` goL l `mappend` goR r + + {-# INLINE foldr #-} + foldr = foldr + {-# INLINE foldl #-} + foldl = foldl + {-# INLINE foldr' #-} + foldr' = foldr' + {-# INLINE foldl' #-} + foldl' = foldl' + #if MIN_VERSION_base(4,8,0) - 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 _ m l r) - | m < 0 = 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 _ m l r) - | m < 0 = 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 #-} + {-# INLINE length #-} + length = size + {-# INLINE null #-} + null = null + {-# INLINE toList #-} + toList = elems -- NB: Foldable.toList only returns values while IntMap.toList returns keys also + {-# INLINE elem #-} + elem = start + where + start _ (IntMap Empty) = False + start v (IntMap (NonEmpty _ minV root)) = v == minV || go v root + + go :: Eq v => v -> Node t v -> Bool + go _ Tip = False + go v (Bin _ boundV l r) = v == boundV || go v l || go v r + {-# INLINABLE sum #-} + sum = foldl' (+) 0 + {-# INLINABLE product #-} + product = foldl' (*) 1 #endif -- | Traverses in order of increasing key. instance Traversable IntMap where - traverse f = traverseWithKey (\_ -> f) {-# INLINE traverse #-} + traverse f = start + where + start (IntMap Empty) = pure (IntMap Empty) + start (IntMap (NonEmpty min minV node)) = liftA2 (\minV' root' -> IntMap (NonEmpty min minV' root')) (f minV) (goL node) + + goL Tip = pure Tip + goL (Bin max maxV l r) = liftA3 (\l' r' v' -> Bin max v' l' r') (goL l) (goR r) (f maxV) + + goR Tip = pure Tip + goR (Bin min minV l r) = liftA3 (Bin min) (f minV) (goL l) (goR r) -instance NFData a => NFData (IntMap a) where - rnf Nil = () - rnf (Tip _ v) = rnf v - rnf (Bin _ _ l r) = rnf l `seq` rnf r -#if __GLASGOW_HASKELL__ +#if MIN_VERSION_base(4,9,0) +-- | @since 0.5.7 +instance Semigroup (IntMap a) where + (<>) = union + stimes = stimesIdempotentMonoid +#endif + +instance Monoid (IntMap a) where + mempty = empty + mappend = union + +#if __GLASGOW_HASKELL__ >= 708 +-- | @since 0.5.6.2 +instance IsList.IsList (IntMap a) where + type Item (IntMap a) = (Key, a) + fromList = fromListLazy + toList = toList +#endif -{-------------------------------------------------------------------- - A Data instance ---------------------------------------------------------------------} +INSTANCE_TYPEABLE1(IntMap) +#if defined(__GLASGOW_HASKELL__) -- 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) + gfoldl f z im = z fromListLazy `f` (toList im) toConstr _ = fromListConstr gunfold k z c = case constrIndex c of - 1 -> k (z fromList) + 1 -> k (z fromListLazy) _ -> error "gunfold" dataTypeOf _ = intMapDataType dataCast1 f = gcast1 f @@ -537,21 +936,49 @@ fromListConstr = mkConstr intMapDataType "fromList" [] Prefix intMapDataType :: DataType intMapDataType = mkDataType "Data.IntMap.Internal.IntMap" [fromListConstr] - #endif -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} +instance NFData a => NFData (IntMap a) where + rnf (IntMap m) = rnf m + +instance NFData a => NFData (IntMap_ t a) where + rnf Empty = () + rnf (NonEmpty _ v root) = rnf v `seq` rnf root + +instance NFData a => NFData (Node t a) where + rnf Tip = () + rnf (Bin _ v l r) = rnf v `seq` rnf l `seq` rnf r + +-- | /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 = findWithDefault (error $ "IntMap.!: key " ++ show k ++ " is not an element of the map") 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 +(\\) = difference + -- | /O(1)/. Is the map empty? -- --- > Data.IntMap.null (empty) == True +-- > Data.IntMap.null empty == True -- > Data.IntMap.null (singleton 1 'a') == False - null :: IntMap a -> Bool -null Nil = True -null _ = False -{-# INLINE null #-} +null (IntMap Empty) = True +null _ = False -- | /O(n)/. Number of elements in the map. -- @@ -559,121 +986,107 @@ null _ = False -- > size (singleton 1 'a') == 1 -- > size (fromList([(1,'a'), (2,'c'), (3,'b')])) == 3 size :: IntMap a -> Int -size = go 0 +size (IntMap Empty) = 0 +size (IntMap (NonEmpty _ _ node)) = sizeNode 0 node where - go !acc (Bin _ _ l r) = go (go acc l) r - go acc (Tip _ _) = 1 + acc - go acc Nil = acc + -- A binary tree will always have exactly one more leaf than it has internal nodes. Counting + -- the leaves is also more efficient since we don't have to add a bunch of zeros. From a + -- different perspective, where the tree is seen as a shuffled PATRICIA tree, the data was + -- "originally" stored in the leaves, so counting those locations gives the size of the map. + sizeNode :: Int -> Node t a -> Int + sizeNode !acc Tip = acc + 1 + sizeNode !acc (Bin _ _ l r) = sizeNode (sizeNode acc l) r + +-- | /O(min(n,W)/. Lookup the value at a key in the map, returning the result +-- to the provided continuations. See also 'lookup'. +-- +-- When 'lookupChurch' is given two arguments (the continuations), it is +-- inlined to the call site. You should therefore use 'lookupChurch' only to +-- define custom lookup functions. +{-# INLINE lookupChurch #-} +lookupChurch :: r -> (a -> r) -> Key -> IntMap a -> r +lookupChurch nothing just = search + where + search !_ (IntMap Empty) = nothing + search !k (IntMap (NonEmpty rootMin rootMinV root)) = case compareMinBound k rootMin of + OutOfBound -> nothing + Matched -> just rootMinV + InBound -> goL (xor k rootMin) root + where + goL !_ Tip = nothing + goL !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goL xorCache l + | otherwise -> goR xorCacheMax r + OutOfBound -> nothing + Matched -> just maxV + where xorCacheMax = xor k max + + goR !_ Tip = nothing + goR !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goR xorCache r + | otherwise -> goL xorCacheMin l + OutOfBound -> nothing + Matched -> just minV + where xorCacheMin = xor k min -- | /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 m l r) | nomatch k p m = False - | zero k m = go l - | otherwise = go r - go (Tip kx _) = k == kx - go Nil = False +member = lookupChurch False (const True) -- | /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 +notMember !k = not . member k -- | /O(min(n,W))/. Lookup the value at a key in the map. See also 'Data.Map.lookup'. - --- See Note: Local 'go' functions and capturing] lookup :: Key -> IntMap a -> Maybe a -lookup !k = go - where - go (Bin p m l r) | nomatch k p m = Nothing - | zero k m = go l - | otherwise = go r - go (Tip kx x) | k == kx = Just x - | otherwise = Nothing - go Nil = Nothing - +lookup = lookupChurch Nothing Just --- See Note: Local 'go' functions and capturing] -find :: Key -> IntMap a -> a -find !k = go - where - go (Bin p m l r) | nomatch k p m = not_found - | zero k m = go l - | otherwise = go r - go (Tip kx x) | k == kx = x - | otherwise = not_found - go Nil = not_found - - not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map") - --- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@ +-- | /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 m l r) | nomatch k p m = def - | zero k m = go l - | otherwise = go r - go (Tip kx x) | k == kx = x - | otherwise = def - go Nil = def +findWithDefault def = lookupChurch def id -- | /O(log n)/. Find largest key smaller than the given one and return the -- corresponding (key, value) pair. -- -- > lookupLT 3 (fromList [(3,'a'), (5,'b')]) == Nothing -- > lookupLT 4 (fromList [(3,'a'), (5,'b')]) == Just (3, 'a') - --- See Note: Local 'go' functions and capturing. lookupLT :: Key -> IntMap a -> Maybe (Key, a) -lookupLT !k t = case t of - Bin _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r - _ -> go Nil t +lookupLT !k = start where - go def (Bin p m l r) - | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r - | zero k m = go def l - | otherwise = go l r - go def (Tip ky y) - | k <= ky = unsafeFindMax def - | otherwise = Just (ky, y) - go def Nil = unsafeFindMax def + start (IntMap Empty) = Nothing + start (IntMap (NonEmpty min minV node)) + | boundKey min >= k = Nothing + | otherwise = Just (goL (xor k min) min minV node) + + goL !_ min minV Tip = (boundKey min, minV) + goL !xorCache min minV (Bin max maxV l r) + | boundKey max < k = (boundKey max, maxV) + | xorCache < xorCacheMax = goL xorCache min minV l + | otherwise = goR xorCacheMax r min minV l + where + xorCacheMax = xor k max --- | /O(log n)/. Find smallest key greater than the given one and return the --- corresponding (key, value) pair. --- --- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') --- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing + goR !_ Tip fMin fMinV fallback = getMax fMin fMinV fallback + goR !xorCache (Bin min minV l r) fMin fMinV fallback + | boundKey min >= k = getMax fMin fMinV fallback + | xorCache < xorCacheMin = goR xorCache r min minV l + | otherwise = goL xorCacheMin min minV l + where + xorCacheMin = xor k min --- See Note: Local 'go' functions and capturing. -lookupGT :: Key -> IntMap a -> Maybe (Key, a) -lookupGT !k t = case t of - Bin _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r - _ -> go Nil t - where - go def (Bin p m l r) - | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def - | zero k m = go r l - | otherwise = go def r - go def (Tip ky y) - | k >= ky = unsafeFindMin def - | otherwise = Just (ky, y) - go def Nil = unsafeFindMin def + getMax min minV Tip = (boundKey min, minV) + getMax _ _ (Bin max maxV _ _) = (boundKey max, maxV) -- | /O(log n)/. Find largest key smaller or equal to the given one and return -- the corresponding (key, value) pair. @@ -681,21 +1094,67 @@ lookupGT !k t = case t of -- > 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 _ m l r | m < 0 -> if k >= 0 then go r l else go Nil r - _ -> go Nil t +lookupLE !k = start where - go def (Bin p m l r) - | nomatch k p m = if k < p then unsafeFindMax def else unsafeFindMax r - | zero k m = go def l - | otherwise = go l r - go def (Tip ky y) - | k < ky = unsafeFindMax def - | otherwise = Just (ky, y) - go def Nil = unsafeFindMax def + start (IntMap Empty) = Nothing + start (IntMap (NonEmpty min minV node)) + | boundKey min > k = Nothing + | otherwise = Just (goL (xor k min) min minV node) + + goL !_ min minV Tip = (boundKey min, minV) + goL !xorCache min minV (Bin max maxV l r) + | boundKey max <= k = (boundKey max, maxV) + | xorCache < xorCacheMax = goL xorCache min minV l + | otherwise = goR xorCacheMax r min minV l + where + xorCacheMax = xor k max + + goR !_ Tip fMin fMinV fallback = getMax fMin fMinV fallback + goR !xorCache (Bin min minV l r) fMin fMinV fallback + | boundKey min > k = getMax fMin fMinV fallback + | xorCache < xorCacheMin = goR xorCache r min minV l + | otherwise = goL xorCacheMin min minV l + where + xorCacheMin = xor k min + + getMax min minV Tip = (boundKey min, minV) + getMax _ _ (Bin max maxV _ _) = (boundKey max, maxV) + +-- | /O(log n)/. Find smallest key greater than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupGT 4 (fromList [(3,'a'), (5,'b')]) == Just (5, 'b') +-- > lookupGT 5 (fromList [(3,'a'), (5,'b')]) == Nothing +lookupGT :: Key -> IntMap a -> Maybe (Key, a) +lookupGT !k = start + where + start (IntMap Empty) = Nothing + start (IntMap (NonEmpty min minV Tip)) + | boundKey min <= k = Nothing + | otherwise = Just (boundKey min, minV) + start (IntMap (NonEmpty min minV (Bin max maxV l r))) + | boundKey max <= k = Nothing + | otherwise = Just (goR (xor k max) max maxV (Bin min minV l r)) + + goL !_ Tip fMax fMaxV fallback = getMin fMax fMaxV fallback + goL !xorCache (Bin max maxV l r) fMax fMaxV fallback + | boundKey max <= k = getMin fMax fMaxV fallback + | xorCache < xorCacheMax = goL xorCache l max maxV r + | otherwise = goR xorCacheMax max maxV r + where + xorCacheMax = xor k max + + goR !_ max maxV Tip = (boundKey max, maxV) + goR !xorCache max maxV (Bin min minV l r) + | boundKey min > k = (boundKey min, minV) + | xorCache < xorCacheMin = goR xorCache max maxV r + | otherwise = goL xorCacheMin l max maxV r + where + xorCacheMin = xor k min + + getMin max maxV Tip = (boundKey max, maxV) + getMin _ _ (Bin min minV _ _) = (boundKey min, minV) -- | /O(log n)/. Find smallest key greater or equal to the given one and return -- the corresponding (key, value) pair. @@ -703,2082 +1162,638 @@ lookupLE !k t = case t of -- > 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 _ m l r | m < 0 -> if k >= 0 then go Nil l else go l r - _ -> go Nil t - where - go def (Bin p m l r) - | nomatch k p m = if k < p then unsafeFindMin l else unsafeFindMin def - | zero k m = go r l - | otherwise = go def r - go def (Tip ky y) - | k > ky = unsafeFindMin def - | otherwise = Just (ky, y) - go def Nil = unsafeFindMin def - - --- Helper function for lookupGE and lookupGT. It assumes that if a Bin node is --- given, it has m > 0. -unsafeFindMin :: IntMap a -> Maybe (Key, a) -unsafeFindMin Nil = Nothing -unsafeFindMin (Tip ky y) = Just (ky, y) -unsafeFindMin (Bin _ _ l _) = unsafeFindMin l - --- Helper function for lookupLE and lookupLT. It assumes that if a Bin node is --- given, it has m > 0. -unsafeFindMax :: IntMap a -> Maybe (Key, a) -unsafeFindMax Nil = Nothing -unsafeFindMax (Tip ky y) = Just (ky, y) -unsafeFindMax (Bin _ _ _ r) = unsafeFindMax r - -{-------------------------------------------------------------------- - Disjoint ---------------------------------------------------------------------} --- | /O(n+m)/. 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 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = disjoint1 - | shorter m2 m1 = disjoint2 - | p1 == p2 = disjoint l1 l2 && disjoint r1 r2 - | otherwise = True +lookupGE !k = start where - disjoint1 | nomatch p2 p1 m1 = True - | zero p2 m1 = disjoint l1 t2 - | otherwise = disjoint r1 t2 - disjoint2 | nomatch p1 p2 m2 = True - | zero p1 m2 = disjoint t1 l2 - | otherwise = disjoint t1 r2 - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} + start (IntMap Empty) = Nothing + start (IntMap (NonEmpty min minV Tip)) + | boundKey min < k = Nothing + | otherwise = Just (boundKey min, minV) + start (IntMap (NonEmpty min minV (Bin max maxV l r))) + | boundKey max < k = Nothing + | otherwise = Just (goR (xor k max) max maxV (Bin min minV l r)) + + goL !_ Tip fMax fMaxV fallback = getMin fMax fMaxV fallback + goL !xorCache (Bin max maxV l r) fMax fMaxV fallback + | boundKey max < k = getMin fMax fMaxV fallback + | xorCache < xorCacheMax = goL xorCache l max maxV r + | otherwise = goR xorCacheMax max maxV r + where + xorCacheMax = xor k max + + goR !_ max maxV Tip = (boundKey max, maxV) + goR !xorCache max maxV (Bin min minV l r) + | boundKey min >= k = (boundKey min, minV) + | xorCache < xorCacheMin = goR xorCache max maxV r + | otherwise = goL xorCacheMin l max maxV r + where + xorCacheMin = xor k min + + getMin max maxV Tip = (boundKey max, maxV) + getMin _ _ (Bin min minV _ _) = (boundKey min, minV) + -- | /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 m l r) - | nomatch k p m = link k (Tip k x) p t - | zero k m = Bin p m (insert k x l) r - | otherwise = Bin p m l (insert k x r) -insert k x t@(Tip ky _) - | k==ky = Tip k x - | otherwise = link k (Tip k x) ky t -insert k x Nil = Tip k x - --- right-biased insertion, used by 'union' --- | /O(min(n,W))/. Insert with a combining function. --- @'insertWith' f key value mp@ --- will insert the pair (key, value) into @mp@ if key does --- not exist in the map. If the key does exist, the function will --- insert @f new_value old_value@. --- --- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] --- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] --- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" - -insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWith f k x t - = insertWithKey (\_ x' y' -> f x' y') k x t - --- | /O(min(n,W))/. Insert with a combining function. --- @'insertWithKey' f key value mp@ --- will insert the pair (key, value) into @mp@ if key does --- not exist in the map. If the key does exist, the function will --- insert @f key new_value old_value@. --- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] --- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] --- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" - -insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWithKey f !k x t@(Bin p m l r) - | nomatch k p m = link k (Tip k x) p t - | zero k m = Bin p m (insertWithKey f k x l) r - | otherwise = Bin p m l (insertWithKey f k x r) -insertWithKey f k x t@(Tip ky y) - | k == ky = Tip k (f k x y) - | otherwise = link k (Tip k x) ky t -insertWithKey _ k x Nil = Tip k x - --- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@) --- is a pair where the first element is equal to (@'lookup' k map@) --- and the second element equal to (@'insertWithKey' f k x map@). --- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) --- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) --- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") --- --- This is how to define @insertLookup@ using @insertLookupWithKey@: --- --- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t --- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) --- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) - -insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) -insertLookupWithKey f !k x t@(Bin p m l r) - | nomatch k p m = (Nothing,link k (Tip k x) p t) - | zero k m = let (found,l') = insertLookupWithKey f k x l - in (found,Bin p m l' r) - | otherwise = let (found,r') = insertLookupWithKey f k x r - in (found,Bin p m l r') -insertLookupWithKey f k x t@(Tip ky y) - | k == ky = (Just y,Tip k (f k x y)) - | otherwise = (Nothing,link k (Tip k x) ky t) -insertLookupWithKey _ k x Nil = (Nothing,Tip k x) - - -{-------------------------------------------------------------------- - 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. +empty = IntMap Empty + +-- | /O(min(n,W))/. Insert with a combining function, generic over strictness. +-- Before inserting a value into the map, @'insertWithEval' eval@ will strictly +-- depend on the resut of calling @eval@ on the value to be inserted. Passing +-- an evaluation function that does no evaluation will result in lazy +-- insertion, while passing a function that actually evaluates its argument +-- will result in strict insertion. +-- +-- When 'insertWithEval' is given one argument (the evaluation function), it is +-- inlined at the call site. 'insertWithEval' is intended for defining +-- @insert@-style functions. +{-# INLINE insertWithEval #-} +insertWithEval :: (a -> ()) -> (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithEval eval = start + where + start _ !k v (IntMap Empty) = eval v `seq` IntMap (NonEmpty (Bound k) v Tip) + start combine !k v (IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> IntMap (NonEmpty min minV (goL combine k v (xor k min) min root)) + OutOfBound -> eval v `seq` IntMap (NonEmpty (Bound k) v (insertMinL (xor k min) min minV root)) + Matched -> let v' = combine v minV + in eval v' `seq` IntMap (NonEmpty (Bound k) v' root) + + -- | Insert a key/value pair into a left node when the key is known to be larger + -- than the minimum bound of the subtree. + goL _ !k v !_ !_ Tip = eval v `seq` Bin (Bound k) v Tip Tip + goL combine !k v !xorCache !min (Bin max maxV l r) = case compareMaxBound k max of + -- In the simple case, we recurse into whichever branch is applicable. + InBound | xorCache < xorCacheMax -> Bin max maxV (goL combine k v xorCache min l) r + | otherwise -> Bin max maxV l (goR combine k v xorCacheMax max r) + + -- If the key is the new maximum, then we have two cases to consider. If + -- the split point between 'min' and 'k' is earlier than the split between + -- 'min' and 'max', then we can just immediately create a new node. Otherwise, + -- we need to push 'max' down into the right branch until it arrives at the + -- correct location. + -- + -- We do the this check by simulating a navigation where 'max' is the key, + -- 'min' is the minimum, and 'k' is the maximum. If 'max' belongs on the + -- left side, then the entire old subtree belongs on the left side. If + -- 'max' belongs on the right side, then we have to push it down. + OutOfBound | xor (boundKey max) min < xorCacheMax -> eval v `seq` Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> eval v `seq` Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + + Matched -> let v' = combine v maxV + in eval v' `seq` Bin max v' l r + where xorCacheMax = xor k max + + goR _ !k v !_ !_ Tip = eval v `seq` Bin (Bound k) v Tip Tip + goR combine !k v !xorCache !max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goR combine k v xorCache max r) + | otherwise -> Bin min minV (goL combine k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> eval v `seq` Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> eval v `seq` Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> let v' = combine v minV + in eval v' `seq` Bin min v' l r + where xorCacheMin = xor k min + +-- Small functions that really ought to be defined in Data.IntMap.Lazy but have +-- to be here for the sake of type class implementations +insertLazy :: Key -> a -> IntMap a -> IntMap a +insertLazy = insertWithEval (const ()) const + +fromListLazy :: [(Key, a)] -> IntMap a +fromListLazy = Data.List.foldl' (\t (k, a) -> insertLazy k a t) empty + +mapLazy :: (a -> b) -> IntMap a -> IntMap b +mapLazy f (IntMap m) = IntMap (mapLazy_ f m) + +mapLazy_ :: (a -> b) -> IntMap_ t a -> IntMap_ t b +mapLazy_ _ Empty = Empty +mapLazy_ f (NonEmpty min minV root) = NonEmpty min (f minV) (mapNodeLazy f root) + +mapNodeLazy :: (a -> b) -> Node t a -> Node t b +mapNodeLazy _ Tip = Tip +mapNodeLazy f (Bin bound value l r) = Bin bound (f value) (mapNodeLazy f l) (mapNodeLazy f r) + +-- These need to be here to avoid a circular dependency between +-- Data.IntMap.Strict and Data.IntMap.Merge.Strict. +mapStrict_ :: (a -> b) -> IntMap_ t a -> IntMap_ t b +mapStrict_ _ Empty = Empty +mapStrict_ f (NonEmpty min minV root) = let !minV' = f minV + in NonEmpty min minV' (mapNodeStrict f root) + +mapNodeStrict :: (a -> b) -> Node t a -> Node t b +mapNodeStrict _ Tip = Tip +mapNodeStrict f (Bin bound value l r) = let !value' = f value + in Bin bound value' (mapNodeStrict f l) (mapNodeStrict f r) + +#if USE_REWRITE_RULES +{-# NOINLINE[1] mapLazy #-} +{-# NOINLINE[1] mapLazy_ #-} +{-# NOINLINE[1] mapNodeLazy #-} +{-# NOINLINE[1] mapStrict_ #-} +{-# NOINLINE[1] mapNodeStrict #-} +{-# RULES +"map/map" forall f g m . mapLazy f (mapLazy g m) = mapLazy (f . g) m +"map_/map_" forall f g m . mapLazy_ f (mapLazy_ g m) = mapLazy_ (f . g) m +"mapS_/mapS_" forall f g m . mapStrict_ f (mapStrict_ g m) = mapStrict_ (\v -> f $! g v) m +"mapS_/map_" forall f g m . mapStrict_ f (mapLazy_ g m) = mapStrict_ (f . g) m +"mapNode/mapNode" forall f g n . mapNodeLazy f (mapNodeLazy g n) = mapNodeLazy (f . g) n +"mapNodeS/mapNodeS" forall f g n . mapNodeStrict f (mapNodeStrict g n) = mapNodeStrict (\v -> f $! g v) n +"mapNodeS/mapNode" forall f g n . mapNodeStrict f (mapNodeLazy g n) = mapNodeStrict (f . g) n + #-} +#if __GLASGOW_HASKELL >= 709 +-- Safe coercions were introduced in 7.8, but did not play well with RULES yes. +{-# RULES +"map/coerce" mapLazy coerce = coerce +"map_/coerce" mapLazy_ coerce = coerce +"mapNode/coerce" mapNodeLazy coerce = coerce + #-} +#endif +#endif + +-- | /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 m l r) - | nomatch k p m = t - | zero k m = binCheckLeft p m (delete k l) r - | otherwise = binCheckRight p m l (delete k r) -delete k t@(Tip ky _) - | k == ky = Nil - | otherwise = t -delete _k Nil = Nil - --- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not --- a member of the map, the original map is returned. --- --- > 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 t@(Bin p m l r) - | nomatch k p m = t - | zero k m = Bin p m (adjustWithKey f k l) r - | otherwise = Bin p m l (adjustWithKey f k r) -adjustWithKey f k t@(Tip ky y) - | k == ky = Tip ky (f k y) - | otherwise = t -adjustWithKey _ _ Nil = Nil - - --- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ --- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is --- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. --- --- > let f x = if x == "a" then Just "new a" else Nothing --- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] --- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a -update f - = updateWithKey (\_ x -> f x) - --- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ --- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is --- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] --- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a -updateWithKey f !k t@(Bin p m l r) - | nomatch k p m = t - | zero k m = binCheckLeft p m (updateWithKey f k l) r - | otherwise = binCheckRight p m l (updateWithKey f k r) -updateWithKey f k t@(Tip ky y) - | k == ky = case (f k y) of - Just y' -> Tip ky y' - Nothing -> Nil - | otherwise = t -updateWithKey _ _ Nil = Nil - --- | /O(min(n,W))/. Lookup and update. --- The function returns original value, if it is updated. --- This is different behavior than 'Data.Map.updateLookupWithKey'. --- Returns the original key value if the map entry is deleted. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")]) --- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) --- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") - -updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a) -updateLookupWithKey f !k t@(Bin p m l r) - | nomatch k p m = (Nothing,t) - | zero k m = let !(found,l') = updateLookupWithKey f k l - in (found,binCheckLeft p m l' r) - | otherwise = let !(found,r') = updateLookupWithKey f k r - in (found,binCheckRight p m l r') -updateLookupWithKey f k t@(Tip ky y) - | k==ky = case (f k y) of - Just y' -> (Just y,Tip ky y') - Nothing -> (Just y,Nil) - | otherwise = (Nothing,t) -updateLookupWithKey _ _ Nil = (Nothing,Nil) - - - --- | /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 m l r) - | nomatch k p m = case f Nothing of - Nothing -> t - Just x -> link k (Tip k x) p t - | zero k m = binCheckLeft p m (alter f k l) r - | otherwise = binCheckRight p m l (alter f k r) -alter f k t@(Tip ky y) - | k==ky = case f (Just y) of - Just x -> Tip ky x - Nothing -> Nil - | otherwise = case f Nothing of - Just x -> link k (Tip k x) ky t - Nothing -> Tip ky y -alter f k Nil = case f Nothing of - Just x -> Tip k x - Nothing -> Nil - --- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at --- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, --- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f --- ('lookup' k m)@. --- --- Example: +delete !_ (IntMap Empty) = IntMap Empty +delete !k m@(IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> IntMap (NonEmpty min minV (deleteL k (xor k min) root)) + OutOfBound -> m + Matched -> IntMap (nodeToMapL root) + +-- | Delete a key from a left node. Takes the XOR of the deleted key and +-- the minimum bound of that node. +-- +-- This would normally be a local method of 'delete', but it can be reused in +-- other places. +deleteL :: Key -> Word -> Node L a -> Node L a +deleteL !_ !_ Tip = Tip +deleteL !k !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (deleteL k xorCache l) r + | otherwise -> Bin max maxV l (deleteR k xorCacheMax r) + OutOfBound -> n + Matched -> extractBinL l r + where xorCacheMax = xor k max + +-- | Delete a key from a right node. Takes the XOR of the deleted key and +-- the maximum bound of that node. +-- +-- This would normally be a local method of 'delete', but it can be reused in +-- other places. +deleteR :: Key -> Word -> Node R a -> Node R a +deleteR !_ !_ Tip = Tip +deleteR !k !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (deleteR k xorCache r) + | otherwise -> Bin min minV (deleteL k xorCacheMin l) r + OutOfBound -> n + Matched -> extractBinR l r + where xorCacheMin = xor k min + +-- | /O(n+m)/. The (left-biased) union of two maps. +-- It prefers the first map when duplicate keys are encountered, +-- i.e. (@'union' == 'unionWith' 'const'@). -- --- @ --- 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 +-- > 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 = start + where + start (IntMap Empty) m2 = m2 + start m1 (IntMap Empty) = m1 + start (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (NonEmpty min1 minV1 (goL2 minV2 min1 root1 min2 root2)) + | min1 > min2 = IntMap (NonEmpty min2 minV2 (goL1 minV1 min1 root1 min2 root2)) + | otherwise = IntMap (NonEmpty min1 minV1 (goLFused min1 root1 root2)) -- we choose min1 arbitrarily, as min1 == min2 + + goL1 minV1 !min1 !n1 !min2 Tip = insertMinL (xor (boundKey min1) min2) min1 minV1 n1 + goL1 minV1 !min1 !n1 !min2 n2@(Bin max2 _ _ _) | boundsDisjoint min1 max2 = unionDisjointL minV1 min2 n2 min1 n1 + goL1 minV1 !min1 Tip !min2 !n2 = goInsertL1 (boundKey min1) minV1 (xor (boundKey min1) min2) min2 n2 + goL1 minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> Bin max2 maxV2 (goL1 minV1 min1 n1 min2 l2) r2 -- we choose min1 arbitrarily - we just need something from tree 1 + | max1 > max2 -> Bin max1 maxV1 l2 (goR2 maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> Bin max2 maxV2 l2 (goR1 maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2) + | otherwise -> Bin max1 maxV1 l2 (goRFused max1 (Bin min1 minV1 l1 r1) r2) -- we choose max1 arbitrarily, as max1 == max2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL1 minV1 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 maxV1 (goL1 minV1 min1 l1 min2 l2) (goRFused max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> Bin max1 maxV1 (goL1 minV1 min1 l1 min2 n2) r1 + + goL2 minV2 !min1 Tip !min2 !n2 = insertMinL (xor (boundKey min2) min1) min2 minV2 n2 + goL2 minV2 !min1 n1@(Bin max1 _ _ _) !min2 !n2 | boundsDisjoint min2 max1 = unionDisjointL minV2 min1 n1 min2 n2 + goL2 minV2 !min1 !n1 !min2 Tip = goInsertL2 (boundKey min2) minV2 (xor (boundKey min2) min1) min1 n1 + goL2 minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> Bin max1 maxV1 (goL2 minV2 min1 l1 min2 n2) r1 -- we choose min2 arbitrarily - we just need something from tree 2 + | max1 > max2 -> Bin max1 maxV1 l1 (goR2 maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | max1 < max2 -> Bin max2 maxV2 l1 (goR1 maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | otherwise -> Bin max1 maxV1 l1 (goRFused max1 r1 (Bin min2 minV2 l2 r2)) -- we choose max1 arbitrarily, as max1 == max2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL2 minV2 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goL2 minV2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 maxV1 (goL2 minV2 min1 l1 min2 l2) (goRFused max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + LT -> Bin max2 maxV2 (goL2 minV2 min1 n1 min2 l2) r2 + + -- 'goLFused' is called instead of 'goL' if the minimums of the two trees are the same + -- Note that because of this property, the trees cannot be disjoint, so we can skip most of the checks in 'goL' + goLFused !_ Tip !n2 = n2 + goLFused !_ !n1 Tip = n1 + goLFused !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> Bin max2 maxV2 (goLFused min n1 l2) r2 + EQ | max1 > max2 -> Bin max1 maxV1 (goLFused min l1 l2) (goR2 maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goLFused min l1 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 maxV1 (goLFused min l1 l2) (goRFused max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> Bin max1 maxV1 (goLFused min l1 n2) r1 + + goR1 maxV1 !max1 !n1 !max2 Tip = insertMaxR (xor (boundKey max1) max2) max1 maxV1 n1 + goR1 maxV1 !max1 !n1 !max2 n2@(Bin min2 _ _ _) | boundsDisjoint min2 max1 = unionDisjointR maxV1 max1 n1 max2 n2 + goR1 maxV1 !max1 Tip !max2 !n2 = goInsertR1 (boundKey max1) maxV1 (xor (boundKey max1) max2) max2 n2 + goR1 maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> Bin min2 minV2 l2 (goR1 maxV1 max1 n1 max2 r2) -- we choose max1 arbitrarily - we just need something from tree 1 + | min1 < min2 -> Bin min1 minV1 (goL2 minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | min1 > min2 -> Bin min2 minV2 (goL1 minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | otherwise -> Bin min1 minV1 (goLFused min1 (Bin max1 maxV1 l1 r1) l2) r2 -- we choose min1 arbitrarily, as min1 == min2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 minV2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> Bin min1 minV1 (goLFused min1 l1 l2) (goR1 maxV1 max1 r1 max2 r2) -- we choose min1 arbitrarily, as min1 == min2 + GT -> Bin min1 minV1 l1 (goR1 maxV1 max1 r1 max2 n2) + + goR2 maxV2 !max1 Tip !max2 !n2 = insertMaxR (xor (boundKey max2) max1) max2 maxV2 n2 + goR2 maxV2 !max1 n1@(Bin min1 _ _ _) !max2 !n2 | boundsDisjoint min1 max2 = unionDisjointR maxV2 max2 n2 max1 n1 + goR2 maxV2 !max1 !n1 !max2 Tip = goInsertR2 (boundKey max2) maxV2 (xor (boundKey max2) max1) max1 n1 + goR2 maxV2 !max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> Bin min1 minV1 l1 (goR2 maxV2 max1 r1 max2 n2) -- we choose max2 arbitrarily - we just need something from tree 2 + | min1 < min2 -> Bin min1 minV1 (goL2 minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | min1 > min2 -> Bin min2 minV2 (goL1 minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | otherwise -> Bin min1 minV1 (goLFused min1 l1 (Bin max2 maxV2 l2 r2)) r1 -- we choose min1 arbitrarily, as min1 == min2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 minV2 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 minV1 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | otherwise -> Bin min1 minV1 (goLFused min1 l1 l2) (goR2 maxV2 max1 r1 max2 r2) -- we choose min1 arbitrarily, as min1 == min2 + LT -> Bin min2 minV2 l2 (goR2 maxV2 max1 n1 max2 r2) + + -- 'goRFused' is called instead of 'goR' if the minimums of the two trees are the same + -- Note that because of this property, the trees cannot be disjoint, so we can skip most of the checks in 'goR' + goRFused !_ Tip n2 = n2 + goRFused !_ n1 Tip = n1 + goRFused !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> Bin min2 minV2 l2 (goRFused max n1 r2) + EQ | min1 < min2 -> Bin min1 minV1 (goL2 minV2 min1 l1 min2 l2) (goRFused max r1 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 minV1 min1 l1 min2 l2) (goRFused max r1 r2) + | otherwise -> Bin min1 minV1 (goLFused min1 l1 l2) (goRFused max r1 r2) -- we choose min1 arbitrarily, as min1 == min2 + GT -> Bin min1 minV1 l1 (goRFused max r1 n2) + + goInsertL1 k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertL1 k v !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goInsertL1 k v xorCache min l) r + | otherwise -> Bin max maxV l (goInsertR1 k v xorCacheMax max r) + OutOfBound | xor (boundKey max) min < xorCacheMax -> Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Bin max v l r + where xorCacheMax = xor k max + + goInsertR1 k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertR1 k v !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goInsertR1 k v xorCache max r) + | otherwise -> Bin min minV (goInsertL1 k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Bin min v l r + where xorCacheMin = xor k min + + goInsertL2 k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertL2 k v !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goInsertL2 k v xorCache min l) r + | otherwise -> Bin max maxV l (goInsertR2 k v xorCacheMax max r) + OutOfBound | xor (boundKey max) min < xorCacheMax -> Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Bin max maxV l r + where xorCacheMax = xor k max + + goInsertR2 k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertR2 k v !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goInsertR2 k v xorCache max r) + | otherwise -> Bin min minV (goInsertL2 k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Bin min minV l r + where xorCacheMin = xor k min + +unionDisjointL :: a -> Bound L -> Node L a -> Bound L -> Node L a -> Node L a +unionDisjointL _ !_ Tip !_ !_ = error "Data.IntMap.unionDisjoint: impossible" +unionDisjointL minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 Tip + | xorBounds min1 max1 < xorBounds min2 max1 = Bin (minToMax min2) minV2 n1 Tip + | otherwise = Bin (minToMax min2) minV2 l1 (insertMaxR (xor (boundKey max1) min2) max1 maxV1 r1) +unionDisjointL minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 (Bin max2 maxV2 l2 r2) + | not (xorBounds min1 max1 `ltMSB` xorBounds min1 max2) = Bin max2 maxV2 l1 (unionDisjointR maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | not (xorBounds min2 max2 `ltMSB` xorBounds min1 max2) = Bin max2 maxV2 (unionDisjointL minV2 min1 n1 min2 l2) r2 + | otherwise = Bin max2 maxV2 n1 (Bin min2 minV2 l2 r2) + +unionDisjointR :: a -> Bound R -> Node R a -> Bound R -> Node R a -> Node R a +unionDisjointR _ !_ !_ !_ Tip = error "Data.IntMap.unionDisjoint: impossible" +unionDisjointR maxV1 !max1 Tip !max2 n2@(Bin min2 minV2 l2 r2) + | xorBounds min2 max2 < xorBounds min2 max1 = Bin (maxToMin max1) maxV1 Tip n2 + | otherwise = Bin (maxToMin max1) maxV1 (insertMinL (xor (boundKey min2) max1) min2 minV2 l2) r2 +unionDisjointR maxV1 !max1 (Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) + | not (xorBounds min2 max2 `ltMSB` xorBounds min1 max2) = Bin min1 minV1 (unionDisjointL minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | not (xorBounds min1 max1 `ltMSB` xorBounds min1 max2) = Bin min1 minV1 l1 (unionDisjointR maxV1 max1 r1 max2 n2) + | otherwise = Bin min1 minV1 (Bin max1 maxV1 l1 r1) n2 -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 :: Data.Foldable.Foldable f => f (IntMap a) -> IntMap a +unions = Data.Foldable.foldl' union empty -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(n+m)/. The (left-biased) union of two maps. --- It prefers the first map when duplicate keys are encountered, --- i.e. (@'union' == 'unionWith' 'const'@). --- --- > union (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "a"), (7, "C")] - -union :: IntMap a -> IntMap a -> IntMap a -union m1 m2 - = mergeWithKey' Bin const id id m1 m2 - --- | /O(n+m)/. The union with a combining function. --- --- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] - -unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. The union with a combining function. --- --- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value --- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] - -unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWithKey f m1 m2 - = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) id id m1 m2 - -{-------------------------------------------------------------------- - Difference ---------------------------------------------------------------------} -- | /O(n+m)/. Difference between two maps (based on keys). -- -- > difference (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 3 "b" - difference :: IntMap a -> IntMap b -> IntMap a -difference m1 m2 - = mergeWithKey (\_ _ _ -> Nothing) id (const Nil) m1 m2 - --- | /O(n+m)/. Difference with a combining function. --- --- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing --- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) --- > == singleton 3 "b:B" - -differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a -differenceWith f m1 m2 - = differenceWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. Difference with a combining function. When two equal keys are --- encountered, the combining function is applied to the key and both values. --- If it returns 'Nothing', the element is discarded (proper set difference). --- If it returns (@'Just' y@), the element is updated with a new value @y@. --- --- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing --- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) --- > == singleton 3 "3:b|B" - -differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a -differenceWithKey f m1 m2 - = mergeWithKey f id (const Nil) m1 m2 - +difference = start + where + start (IntMap Empty) !_ = IntMap Empty + start !m (IntMap Empty) = m + start (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 _ root2)) + | min1 < min2 = IntMap (NonEmpty min1 minV1 (goL2 min1 root1 min2 root2)) + | min1 > min2 = IntMap (goL1 minV1 min1 root1 min2 root2) + | otherwise = IntMap (goLFused min1 root1 root2) + + goL1 minV1 !min1 Tip !min2 !n2 = goLookupL (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL1 minV1 !min1 !n1 !_ Tip = NonEmpty min1 minV1 n1 + goL1 minV1 !min1 n1@(Bin _ _ _ _) !_ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = NonEmpty min1 minV1 n1 + goL1 minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL1 minV1 min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> r2lMap $ NonEmpty max1 maxV1 (goR2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> r2lMap $ goR1 maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + | otherwise -> r2lMap $ goRFused max1 (Bin min1 minV1 l1 r1) r2 + EQ | max1 > max2 -> binL (goL1 minV1 min1 l1 min2 l2) (NonEmpty max1 maxV1 (goR2 max1 r1 max2 r2)) + | max1 < max2 -> binL (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binL (goL1 minV1 min1 l1 min2 l2) (goRFused max1 r1 r2) + GT -> binL (goL1 minV1 min1 l1 min2 n2) (NonEmpty max1 maxV1 r1) + + goL2 !_ Tip !_ !_ = Tip + goL2 !min1 !n1 !min2 !Tip = deleteL (boundKey min2) (xor (boundKey min2) min1) n1 + goL2 !_ n1@(Bin max1 _ _ _) !min2 (Bin _ _ _ _) | boundsDisjoint min2 max1 = n1 + goL2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goL2 min1 n1 min2 l2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL2 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | max1 < max2 -> binNodeMapL (goL2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binNodeMapL (goL2 min1 l1 min2 l2) (goRFused max1 r1 r2) + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> Bin max1 maxV1 (goL2 min1 l1 min2 n2) r1 -- min2 is arbitrary here - we just need something from tree 2 + | max1 > max2 -> Bin max1 maxV1 l1 (goR2 max1 r1 max2 (Bin min2 dummyV l2 r2)) + | max1 < max2 -> binNodeMapL l1 (goR1 maxV1 max1 r1 max2 (Bin min2 dummyV l2 r2)) + | otherwise -> binNodeMapL l1 (goRFused max1 r1 (Bin min2 dummyV l2 r2)) + + goLFused !_ Tip !_ = Empty + goLFused !_ (Bin max1 maxV1 l1 r1) Tip = case deleteMinL max1 maxV1 l1 r1 of + NE min' minV' n' -> NonEmpty min' minV' n' + goLFused !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFused min n1 l2 + EQ | max1 > max2 -> binL (goLFused min l1 l2) (NonEmpty max1 maxV1 (goR2 max1 r1 max2 r2)) + | max1 < max2 -> binL (goLFused min l1 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binL (goLFused min l1 l2) (goRFused max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> binL (goLFused min l1 n2) (NonEmpty max1 maxV1 r1) + + goR1 maxV1 !max1 Tip !max2 !n2 = goLookupR (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR1 maxV1 !max1 !n1 !_ Tip = NonEmpty max1 maxV1 n1 + goR1 maxV1 !max1 n1@(Bin _ _ _ _) !_ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = NonEmpty max1 maxV1 n1 + goR1 maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR1 maxV1 max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> l2rMap $ NonEmpty min1 minV1 (goL2 min1 (Bin max1 maxV1 l1 r1) min2 l2) + | min1 > min2 -> l2rMap $ goL1 minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | otherwise -> l2rMap $ goLFused min1 (Bin max1 maxV1 l1 r1) l2 + EQ | min1 < min2 -> binR (NonEmpty min1 minV1 (goL2 min1 l1 min2 l2)) (goR1 maxV1 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binR (goLFused min1 l1 l2) (goR1 maxV1 max1 r1 max2 r2) + GT -> binR (NonEmpty min1 minV1 l1) (goR1 maxV1 max1 r1 max2 n2) + + goR2 !_ Tip !_ !_ = Tip + goR2 !max1 !n1 !max2 Tip = deleteR (boundKey max2) (xor (boundKey max2) max1) n1 + goR2 !_ n1@(Bin min1 _ _ _) !max2 (Bin _ _ _ _) | boundsDisjoint min1 max2 = n1 + goR2 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goR2 max1 n1 max2 r2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | min1 > min2 -> binMapNodeR (goL1 minV1 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | otherwise -> binMapNodeR (goLFused min1 l1 l2) (goR2 max1 r1 max2 r2) + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> Bin min1 minV1 l1 (goR2 max1 r1 max2 n2) -- max2 is arbitrary here - we just need something from tree 2 + | min1 < min2 -> Bin min1 minV1 (goL2 min1 l1 min2 (Bin max2 dummyV l2 r2)) r1 + | min1 > min2 -> binMapNodeR (goL1 minV1 min1 l1 min2 (Bin max2 dummyV l2 r2)) r1 + | otherwise -> binMapNodeR (goLFused min1 l1 (Bin max2 dummyV l2 r2)) r1 + + goRFused !_ Tip !_ = Empty + goRFused !_ (Bin min1 minV1 l1 r1) Tip = case deleteMaxR min1 minV1 l1 r1 of + NE max' maxV' n' -> NonEmpty max' maxV' n' + goRFused !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFused max n1 r2 + EQ | min1 < min2 -> binR (NonEmpty min1 minV1 (goL2 min1 l1 min2 l2)) (goRFused max r1 r2) + | min1 > min2 -> binR (goL1 minV1 min1 l1 min2 l2) (goRFused max r1 r2) + | otherwise -> binR (goLFused min1 l1 l2) (goRFused max r1 r2) -- we choose min1 arbitrarily, as min1 == min2 + GT -> binR (NonEmpty min1 minV1 l1) (goRFused max r1 n2) + + goLookupL !k v !_ Tip = NonEmpty (Bound k) v Tip + goLookupL !k v !xorCache (Bin max _ l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL k v xorCache l + | otherwise -> goLookupR k v xorCacheMax r + OutOfBound -> NonEmpty (Bound k) v Tip + Matched -> Empty + where xorCacheMax = xor k max + + goLookupR !k v !_ Tip = NonEmpty (Bound k) v Tip + goLookupR !k v !xorCache (Bin min _ l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR k v xorCache r + | otherwise -> goLookupL k v xorCacheMin l + OutOfBound -> NonEmpty (Bound k) v Tip + Matched -> Empty + where xorCacheMin = xor k min + + dummyV = error "impossible" --- TODO(wrengr): re-verify that asymptotic bound --- | /O(n+m)/. 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 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2) - | shorter m1 m2 = difference1 - | shorter m2 m1 = difference2 - | p1 == p2 = bin p1 m1 (withoutKeys l1 l2) (withoutKeys r1 r2) - | otherwise = t1 - where - difference1 - | nomatch p2 p1 m1 = t1 - | zero p2 m1 = binCheckLeft p1 m1 (withoutKeys l1 t2) r1 - | otherwise = binCheckRight p1 m1 l1 (withoutKeys r1 t2) - difference2 - | nomatch p1 p2 m2 = t1 - | zero p1 m2 = withoutKeys t1 l2 - | otherwise = withoutKeys t1 r2 -withoutKeys t1@(Bin p1 m1 _ _) (IntSet.Tip p2 bm2) = - let minbit = bitmapOf p1 - lt_minbit = minbit - 1 - maxbit = bitmapOf (p1 .|. (m1 .|. (m1 - 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 m l r) f - | m .&. IntSet.suffixBitMask /= 0 = - if p .&. IntSet.prefixBitMask == kp then f t else t - | nomatch kp p m = t - | zero kp m = binCheckLeft p m (updatePrefix kp l f) r - | otherwise = binCheckRight p m 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 m l r) = - let leftBits = bitmapOf (p .|. m) - 1 - bmL = bm .&. leftBits - bmR = bm `xor` bmL -- = (bm .&. complement leftBits) - in bin p m (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(n+m)/. The (left-biased) intersection of two maps (based on keys). -- -- > intersection (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "a" - intersection :: IntMap a -> IntMap b -> IntMap a -intersection m1 m2 - = mergeWithKey' bin const (const Nil) (const Nil) m1 m2 - - --- TODO(wrengr): re-verify that asymptotic bound --- | /O(n+m)/. The restriction of a map to the keys in a set. --- --- @ --- m \`restrictKeys\` s = 'filterWithKey' (\k _ -> k ``IntSet.member`` s) m --- @ --- --- @since 0.5.8 -restrictKeys :: IntMap a -> IntSet.IntSet -> IntMap a -restrictKeys t1@(Bin p1 m1 l1 r1) t2@(IntSet.Bin p2 m2 l2 r2) - | shorter m1 m2 = intersection1 - | shorter m2 m1 = intersection2 - | p1 == p2 = bin p1 m1 (restrictKeys l1 l2) (restrictKeys r1 r2) - | otherwise = Nil - where - intersection1 - | nomatch p2 p1 m1 = Nil - | zero p2 m1 = restrictKeys l1 t2 - | otherwise = restrictKeys r1 t2 - intersection2 - | nomatch p1 p2 m2 = Nil - | zero p1 m2 = restrictKeys t1 l2 - | otherwise = restrictKeys t1 r2 -restrictKeys t1@(Bin p1 m1 _ _) (IntSet.Tip p2 bm2) = - let minbit = bitmapOf p1 - ge_minbit = complement (minbit - 1) - maxbit = bitmapOf (p1 .|. (m1 .|. (m1 - 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 m l r) - | m .&. IntSet.suffixBitMask /= 0 = - if p .&. IntSet.prefixBitMask == kp then t else Nil - | nomatch kp p m = Nil - | zero kp m = 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 m l r) = - let leftBits = bitmapOf (p .|. m) - 1 - bmL = bm .&. leftBits - bmR = bm `xor` bmL -- = (bm .&. complement leftBits) - in bin p m (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(n+m)/. The intersection with a combining function. --- --- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" - -intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c -intersectionWith f m1 m2 - = intersectionWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. The intersection with a combining function. --- --- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar --- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" - -intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c -intersectionWithKey f m1 m2 - = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 (f k1 x1 x2)) (const Nil) (const Nil) m1 m2 - -{-------------------------------------------------------------------- - MergeWithKey ---------------------------------------------------------------------} - --- | /O(n+m)/. A high-performance universal combining function. Using --- 'mergeWithKey', all combining functions can be defined without any loss of --- efficiency (with exception of 'union', 'difference' and 'intersection', --- where sharing of some nodes is lost with 'mergeWithKey'). --- --- Please make sure you know what is going on when using 'mergeWithKey', --- otherwise you can be surprised by unexpected code growth or even --- corruption of the data structure. --- --- When 'mergeWithKey' is given three arguments, it is inlined to the call --- site. You should therefore use 'mergeWithKey' only to define your custom --- combining functions. For example, you could define 'unionWithKey', --- 'differenceWithKey' and 'intersectionWithKey' as --- --- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 --- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 --- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 --- --- When calling @'mergeWithKey' combine only1 only2@, a function combining two --- 'IntMap's is created, such that --- --- * if a key is present in both maps, it is passed with both corresponding --- values to the @combine@ function. Depending on the result, the key is either --- present in the result with specified value, or is left out; --- --- * a nonempty subtree present only in the first map is passed to @only1@ and --- the output is added to the result; --- --- * a nonempty subtree present only in the second map is passed to @only2@ and --- the output is added to the result. --- --- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. --- The values can be modified arbitrarily. Most common variants of @only1@ and --- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or --- @'filterWithKey' f@ could be used for any @f@. - -mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) - -> IntMap a -> IntMap b -> IntMap c -mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2 - where -- We use the lambda form to avoid non-exhaustive pattern matches warning. - combine = \(Tip k1 x1) (Tip _k2 x2) -> - case f k1 x1 x2 of - Nothing -> Nil - Just x -> Tip k1 x - {-# INLINE combine #-} -{-# INLINE mergeWithKey #-} - --- Slightly more general version of mergeWithKey. It differs in the following: --- --- * the combining function operates on maps instead of keys and values. The --- reason is to enable sharing in union, difference and intersection. --- --- * mergeWithKey' is given an equivalent of bin. The reason is that in union*, --- Bin constructor can be used, because we know both subtrees are nonempty. - -mergeWithKey' :: (Prefix -> Mask -> IntMap c -> IntMap c -> IntMap c) - -> (IntMap a -> IntMap b -> IntMap c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) - -> IntMap a -> IntMap b -> IntMap c -mergeWithKey' bin' f g1 g2 = go - where - go t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = merge1 - | shorter m2 m1 = merge2 - | p1 == p2 = bin' p1 m1 (go l1 l2) (go r1 r2) - | otherwise = maybe_link p1 (g1 t1) p2 (g2 t2) - where - merge1 | nomatch p2 p1 m1 = maybe_link p1 (g1 t1) p2 (g2 t2) - | zero p2 m1 = bin' p1 m1 (go l1 t2) (g1 r1) - | otherwise = bin' p1 m1 (g1 l1) (go r1 t2) - merge2 | nomatch p1 p2 m2 = maybe_link p1 (g1 t1) p2 (g2 t2) - | zero p1 m2 = bin' p2 m2 (go t1 l2) (g2 r2) - | otherwise = bin' p2 m2 (g2 l2) (go t1 r2) - - go t1'@(Bin _ _ _ _) t2'@(Tip k2' _) = merge0 t2' k2' t1' - where - merge0 t2 k2 t1@(Bin p1 m1 l1 r1) - | nomatch k2 p1 m1 = maybe_link p1 (g1 t1) k2 (g2 t2) - | zero k2 m1 = bin' p1 m1 (merge0 t2 k2 l1) (g1 r1) - | otherwise = bin' p1 m1 (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 m2 l2 r2) - | nomatch k1 p2 m2 = maybe_link k1 (g1 t1) p2 (g2 t2) - | zero k1 m2 = bin' p2 m2 (merge0 t1 k1 l2) (g2 r2) - | otherwise = bin' p2 m2 (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 t2 = g2 t2 - - maybe_link _ Nil _ t2 = t2 - maybe_link _ t1 _ Nil = t1 - maybe_link p1 t1 p2 t2 = link p1 t1 p2 t2 - {-# INLINE maybe_link #-} -{-# INLINE mergeWithKey' #-} - - -{-------------------------------------------------------------------- - 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) +intersection = start 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 -#if !MIN_VERSION_base(4,8,0) - return = pure -#endif - m >>= f = - traverseMaybeMissing $ \k x -> do - res1 <- missingKey m k x - case res1 of - Nothing -> pure Nothing - Just r -> missingKey (f r) k x - {-# INLINE (>>=) #-} - + start (IntMap Empty) !_ = IntMap Empty + start !_ (IntMap Empty) = IntMap Empty + start (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 _ root2)) + | min1 < min2 = IntMap (goL2 min1 root1 min2 root2) + | min1 > min2 = IntMap (goL1 minV1 min1 root1 min2 root2) + | otherwise = IntMap (NonEmpty min1 minV1 (goLFused min1 root1 root2)) -- we choose min1 arbitrarily, as min1 == min2 + + -- TODO: This scheme might produce lots of unnecessary l2r and r2l calls. This should be rectified. + + goL1 _ !_ !_ !_ Tip = Empty + goL1 minV1 min1 Tip min2 n2 = goLookupL1 (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL1 _ min1 (Bin _ _ _ _) _ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = Empty + goL1 minV1 min1 n1@(Bin max1 maxV1 l1 r1) min2 n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL1 minV1 min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> r2lMap $ goR2 max1 (Bin min1 minV1 l1 r1) max2 r2 + | max1 < max2 -> r2lMap $ goR1 maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + | otherwise -> r2lMap $ NonEmpty max1 maxV1 (goRFused max1 (Bin min1 minV1 l1 r1) r2) + EQ | max1 > max2 -> binL (goL1 minV1 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | max1 < max2 -> binL (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binL (goL1 minV1 min1 l1 min2 l2) (NonEmpty max1 maxV1 (goRFused max1 r1 r2)) + GT -> goL1 minV1 min1 l1 min2 n2 + + goL2 !_ Tip !_ !_ = Empty + goL2 min1 n1 min2 Tip = goLookupL2 (boundKey min2) (xor (boundKey min2) min1) n1 + goL2 _ (Bin max1 _ _ _) min2 (Bin _ _ _ _) | boundsDisjoint min2 max1 = Empty + goL2 min1 n1@(Bin max1 maxV1 l1 r1) min2 n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goL2 min1 n1 min2 l2 + EQ | max1 > max2 -> binL (goL2 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | max1 < max2 -> binL (goL2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binL (goL2 min1 l1 min2 l2) (NonEmpty max1 maxV1 (goRFused max1 r1 r2)) + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> goL2 min1 l1 min2 n2 -- min2 is arbitrary here - we just need something from tree 2 + | max1 > max2 -> r2lMap $ goR2 max1 r1 max2 (Bin min2 dummyV l2 r2) + | max1 < max2 -> r2lMap $ goR1 maxV1 max1 r1 max2 (Bin min2 dummyV l2 r2) + | otherwise -> r2lMap $ NonEmpty max1 maxV1 (goRFused max1 r1 (Bin min2 dummyV l2 r2)) + + goLFused !_ Tip !_ = Tip + goLFused !_ !_ Tip = Tip + goLFused !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFused min n1 l2 + EQ | max1 > max2 -> binNodeMapL (goLFused min l1 l2) (goR2 max1 r1 max2 r2) + | max1 < max2 -> binNodeMapL (goLFused min l1 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 maxV1 (goLFused min l1 l2) (goRFused max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> goLFused min l1 n2 + + goR1 _ !_ !_ !_ Tip = Empty + goR1 maxV1 max1 Tip max2 n2 = goLookupR1 (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR1 _ max1 (Bin _ _ _ _) _ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = Empty + goR1 maxV1 max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR1 maxV1 max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> l2rMap $ goL2 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | min1 > min2 -> l2rMap $ goL1 minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | otherwise -> l2rMap $ NonEmpty min1 minV1 (goLFused min1 (Bin max1 maxV1 l1 r1) l2) + EQ | min1 < min2 -> binR (goL2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> binR (NonEmpty min1 minV1 (goLFused min1 l1 l2)) (goR1 maxV1 max1 r1 max2 r2) + GT -> goR1 maxV1 max1 r1 max2 n2 + + goR2 !_ Tip !_ !_ = Empty + goR2 max1 n1 max2 Tip = goLookupR2 (boundKey max2) (xor (boundKey max2) max1) n1 + goR2 _ (Bin min1 _ _ _) max2 (Bin _ _ _ _) | boundsDisjoint min1 max2 = Empty + goR2 max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goR2 max1 n1 max2 r2 + EQ | min1 < min2 -> binR (goL2 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 minV1 min1 l1 min2 l2) (goR2 max1 r1 max2 r2) + | otherwise -> binR (NonEmpty min1 minV1 (goLFused min1 l1 l2)) (goR2 max1 r1 max2 r2) + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> goR2 max1 r1 max2 n2 -- max2 is arbitrary here - we just need something from tree 2 + | min1 < min2 -> l2rMap $ goL2 min1 l1 min2 (Bin max2 dummyV l2 r2) + | min1 > min2 -> l2rMap $ goL1 minV1 min1 l1 min2 (Bin max2 dummyV l2 r2) + | otherwise -> l2rMap $ NonEmpty min1 minV1 (goLFused min1 l1 (Bin max2 dummyV l2 r2)) + + goRFused !_ Tip !_ = Tip + goRFused !_ !_ Tip = Tip + goRFused !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFused max n1 r2 + EQ | min1 < min2 -> binMapNodeR (goL2 min1 l1 min2 l2) (goRFused max r1 r2) + | min1 > min2 -> binMapNodeR (goL1 minV1 min1 l1 min2 l2) (goRFused max r1 r2) + | otherwise -> Bin min1 minV1 (goLFused min1 l1 l2) (goRFused max r1 r2) -- we choose min1 arbitrarily, as min1 == min2 + GT -> goRFused max r1 n2 + + goLookupL1 !_ _ !_ Tip = Empty + goLookupL1 k v !xorCache (Bin max _ l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL1 k v xorCache l + | otherwise -> goLookupR1 k v xorCacheMax r + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) v Tip + where xorCacheMax = xor k max + + goLookupR1 !_ _ !_ Tip = Empty + goLookupR1 k v !xorCache (Bin min _ l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR1 k v xorCache r + | otherwise -> goLookupL1 k v xorCacheMin l + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) v Tip + where xorCacheMin = xor k min + + goLookupL2 !_ !_ Tip = Empty + goLookupL2 k !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL2 k xorCache l + | otherwise -> goLookupR2 k xorCacheMax r + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) maxV Tip + where xorCacheMax = xor k max + + goLookupR2 !_ !_ Tip = Empty + goLookupR2 k !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR2 k xorCache r + | otherwise -> goLookupL2 k xorCacheMin l + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) minV Tip + where xorCacheMin = xor k min + + dummyV = error "impossible" --- | 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 #-} - - -#if !MIN_VERSION_base(4,8,0) -newtype Identity a = Identity {runIdentity :: a} - -instance Functor Identity where - fmap f (Identity x) = Identity (f x) - -instance Applicative Identity where - pure = Identity - Identity f <*> Identity x = Identity (f x) -#endif - --- | 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 -#if !MIN_VERSION_base(4,8,0) - return = pure -#endif - m >>= f = - zipWithMaybeAMatched $ \k x y -> do - res <- runWhenMatched m k x y - case res of - Nothing -> pure Nothing - Just r -> runWhenMatched (f r) k x y - {-# INLINE (>>=) #-} - - --- | Map covariantly over a @'WhenMatched' f 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 m l r) - | m < 0 = liftA2 (flip (bin p m)) (filterWithKeyA f r) (filterWithKeyA f l) - | otherwise = liftA2 (bin p m) (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. -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 m l r) - | m < 0 = liftA2 (flip (bin p m)) (go r) (go l) - | otherwise = liftA2 (bin p m) (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 m2 l2 r2) - | nomatch k1 p2 m2 = linkA k1 (subsingletonBy g1k k1 x1) p2 (g2t t2) - | zero k1 m2 = binA p2 m2 (merge2 l2) (g2t r2) - | otherwise = binA p2 m2 (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 m1 l1 r1) - | nomatch k2 p1 m1 = linkA p1 (g1t t1) k2 (subsingletonBy g2k k2 x2) - | zero k2 m1 = binA p1 m1 (merge1 l1) (g1t r1) - | otherwise = binA p1 m1 (g1t l1) (merge1 r1) - merge1 (Tip k1 x1) = mergeTips k1 x1 k2 x2 - merge1 Nil = subsingletonBy g2k k2 x2 - - go t1@(Bin p1 m1 l1 r1) t2@(Bin p2 m2 l2 r2) - | shorter m1 m2 = merge1 - | shorter m2 m1 = merge2 - | p1 == p2 = binA p1 m1 (go l1 l2) (go r1 r2) - | otherwise = linkA p1 (g1t t1) p2 (g2t t2) - where - merge1 | nomatch p2 p1 m1 = linkA p1 (g1t t1) p2 (g2t t2) - | zero p2 m1 = binA p1 m1 (go l1 t2) (g1t r1) - | otherwise = binA p1 m1 (g1t l1) (go r1 t2) - merge2 | nomatch p1 p2 m2 = linkA p1 (g1t t1) p2 (g2t t2) - | zero p1 m2 = binA p2 m2 (go t1 l2) (g2t r2) - | otherwise = binA p2 m2 (g2t l2) (go t1 r2) - - 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 - => Prefix -> f (IntMap a) - -> Prefix -> f (IntMap a) - -> f (IntMap a) - linkA p1 t1 p2 t2 - | zero p1 m = binA p m t1 t2 - | otherwise = binA p m t2 t1 - where - m = branchMask p1 p2 - p = mask p1 m - {-# INLINE linkA #-} - - -- A variant of 'bin' that ensures that effects for negative keys are executed - -- first. - binA - :: Applicative f - => Prefix - -> Mask - -> f (IntMap a) - -> f (IntMap a) - -> f (IntMap a) - binA p m a b - | m < 0 = liftA2 (flip (bin p m)) b a - | otherwise = liftA2 (bin p m) 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 m l r | m < 0 -> binCheckRight p m l (go f r) - _ -> go f t - where - go f' (Bin p m l r) = binCheckLeft p m (go f' l) r - go f' (Tip k y) = case f' k y of - Just y' -> Tip k y' - Nothing -> Nil - go _ Nil = error "updateMinWithKey Nil" - --- | /O(min(n,W))/. Update the value at the maximal key. --- --- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")] --- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" - -updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a -updateMaxWithKey f t = - case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r - _ -> go f t - where - go f' (Bin p m l r) = binCheckRight p m l (go f' r) - go f' (Tip k y) = case f' k y of - Just y' -> Tip k y' - Nothing -> Nil - go _ Nil = error "updateMaxWithKey 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 m l r | m < 0 -> - case go l of View k a l' -> View k a (binCheckLeft p m l' r) - _ -> go t - where - go (Bin p m l r) = - case go r of View k a r' -> View k a (binCheckRight p m 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 m l r | m < 0 -> - case go r of - View k a r' -> View k a (binCheckRight p m l r') - _ -> go t - where - go (Bin p m l r) = - case go l of View k a l' -> View k a (binCheckLeft p m 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 - --- | /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 _ m l r) - | m < 0 = go r - | otherwise = go l - where go (Tip k v) = Just (k,v) - go (Bin _ _ l' _) = go l' - go Nil = Nothing - --- | /O(min(n,W))/. The minimal key of the map. Calls 'error' if the map is empty. --- Use 'minViewWithKey' if the map may be empty. -findMin :: IntMap a -> (Key, a) -findMin t - | Just r <- lookupMin t = r - | otherwise = error "findMin: empty map has no minimal element" - --- | /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 _ m l r) - | m < 0 = go l - | otherwise = go r - where go (Tip k v) = Just (k,v) - go (Bin _ _ _ r') = go r' - go Nil = Nothing - --- | /O(min(n,W))/. The maximal key of the map. Calls 'error' if the map is empty. --- Use 'maxViewWithKey' if the map may be 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(n+m)/. Is this a proper submap? (ie. a submap but not equal). --- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). -isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool -isProperSubmapOf m1 m2 - = isProperSubmapOfBy (==) m1 m2 - -{- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). - The expression (@'isProperSubmapOfBy' f m1 m2@) returns 'True' when - @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 m1 l1 r1) (Bin p2 m2 l2 r2) - | shorter m1 m2 = GT - | shorter m2 m1 = submapCmpLt - | p1 == p2 = submapCmpEq - | otherwise = GT -- disjoint - where - submapCmpLt | nomatch p1 p2 m2 = GT - | zero p1 m2 = submapCmp predicate t1 l2 - | otherwise = submapCmp predicate t1 r2 - submapCmpEq = case (submapCmp predicate l1 l2, submapCmp predicate r1 r2) of - (GT,_ ) -> GT - (_ ,GT) -> GT - (EQ,EQ) -> EQ - _ -> LT - -submapCmp _ (Bin _ _ _ _) _ = GT -submapCmp predicate (Tip kx x) (Tip ky y) - | (kx == ky) && predicate x y = EQ - | otherwise = GT -- disjoint -submapCmp predicate (Tip k x) t - = case lookup k t of - Just y | predicate x y -> LT - _ -> GT -- disjoint -submapCmp _ Nil Nil = EQ -submapCmp _ Nil _ = LT - --- | /O(n+m)/. Is this a submap? --- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). -isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool -isSubmapOf m1 m2 - = isSubmapOfBy (==) m1 m2 - -{- | /O(n+m)/. - The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if - all keys in @m1@ are in @m2@, and when @f@ returns 'True' when - applied to their respective values. For example, the following - expressions are all 'True': - - > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) - - But the following are all 'False': - - > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) - > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) --} -isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool -isSubmapOfBy predicate t1@(Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - | shorter m1 m2 = False - | shorter m2 m1 = match p1 p2 m2 && - if zero p1 m2 - then isSubmapOfBy predicate t1 l2 - else isSubmapOfBy predicate t1 r2 - | otherwise = (p1==p2) && isSubmapOfBy predicate l1 l2 && isSubmapOfBy predicate r1 r2 -isSubmapOfBy _ (Bin _ _ _ _) _ = False -isSubmapOfBy predicate (Tip k x) t = case lookup k t of - Just y -> predicate x y - Nothing -> False -isSubmapOfBy _ Nil _ = True - -{-------------------------------------------------------------------- - Mapping ---------------------------------------------------------------------} --- | /O(n)/. Map a function over all values in the map. --- --- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] - -map :: (a -> b) -> IntMap a -> IntMap b -map f = go - where - go (Bin p m l r) = Bin p m (go l) (go r) - go (Tip k x) = Tip k (f x) - go Nil = Nil - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] map #-} -{-# RULES -"map/map" forall f g xs . map f (map g xs) = map (f . g) xs - #-} -#endif -#if __GLASGOW_HASKELL__ >= 709 --- Safe coercions were introduced in 7.8, but did not play well with RULES yet. -{-# RULES -"map/coerce" map coerce = coerce - #-} -#endif - --- | /O(n)/. Map a function over all values in the map. --- --- > let f key x = (show key) ++ ":" ++ x --- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")] - -mapWithKey :: (Key -> a -> b) -> IntMap a -> IntMap b -mapWithKey f t - = case t of - Bin p m l r -> Bin p m (mapWithKey f l) (mapWithKey f r) - Tip k x -> Tip k (f k x) - Nil -> Nil - -#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 m l r) - | m < 0 = liftA2 (flip (Bin p m)) (go r) (go l) - | otherwise = liftA2 (Bin p m) (go l) (go r) -{-# INLINE traverseWithKey #-} - --- | /O(n)/. The function @'mapAccum'@ threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a b = (a ++ b, b ++ "X") --- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) - -mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccum f = mapAccumWithKey (\a' _ x -> f a' x) - --- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") --- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) - -mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccumWithKey f a t - = mapAccumL f a t - --- | /O(n)/. The function @'mapAccumL'@ threads an accumulating --- argument through the map in ascending order of keys. -mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccumL f a t - = case t of - Bin p m l r - | m < 0 -> - let (a1,r') = mapAccumL f a r - (a2,l') = mapAccumL f a1 l - in (a2,Bin p m l' r') - | otherwise -> - let (a1,l') = mapAccumL f a l - (a2,r') = mapAccumL f a1 r - in (a2,Bin p m l' r') - Tip k x -> let (a',x') = f a k x in (a',Tip k x') - Nil -> (a,Nil) - --- | /O(n)/. The function @'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 m l r - | m < 0 -> - let (a1,l') = mapAccumRWithKey f a l - (a2,r') = mapAccumRWithKey f a1 r - in (a2,Bin p m l' r') - | otherwise -> - let (a1,r') = mapAccumRWithKey f a r - (a2,l') = mapAccumRWithKey f a1 l - in (a2,Bin p m l' r') - Tip k x -> let (a',x') = f a k x in (a',Tip k x') - Nil -> (a,Nil) - --- | /O(n*min(n,W))/. --- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. --- --- The size of the result may be smaller if @f@ maps two or more distinct --- keys to the same new key. In this case the value at the greatest of the --- original keys is retained. --- --- > mapKeys (+ 1) (fromList [(5,"a"), (3,"b")]) == fromList [(4, "b"), (6, "a")] --- > mapKeys (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "c" --- > mapKeys (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "c" - -mapKeys :: (Key->Key) -> IntMap a -> IntMap a -mapKeys f = fromList . foldrWithKey (\k x xs -> (f k, x) : xs) [] - --- | /O(n*min(n,W))/. --- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. --- --- The size of the result may be smaller if @f@ maps two or more distinct --- keys to the same new key. In this case the associated values will be --- combined using @c@. --- --- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" --- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" - -mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a -mapKeysWith c f - = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] - --- | /O(n*min(n,W))/. --- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ --- is strictly monotonic. --- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. --- /The precondition is not checked./ --- Semi-formally, we have: --- --- > and [x < y ==> f x < f y | x <- ls, y <- ls] --- > ==> mapKeysMonotonic f s == mapKeys f s --- > where ls = keys s --- --- This means that @f@ maps distinct original keys to distinct resulting keys. --- This function has slightly better performance than 'mapKeys'. --- --- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] - -mapKeysMonotonic :: (Key->Key) -> IntMap a -> IntMap a -mapKeysMonotonic f - = fromDistinctAscList . foldrWithKey (\k x xs -> (f k, x) : xs) [] - -{-------------------------------------------------------------------- - Filter ---------------------------------------------------------------------} --- | /O(n)/. Filter all values that satisfy some predicate. --- --- > filter (> "a") (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" --- > filter (> "x") (fromList [(5,"a"), (3,"b")]) == empty --- > filter (< "a") (fromList [(5,"a"), (3,"b")]) == empty - -filter :: (a -> Bool) -> IntMap a -> IntMap a -filter p m - = filterWithKey (\_ x -> p x) m - --- | /O(n)/. Filter all keys\/values that satisfy some predicate. --- --- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a -filterWithKey predicate = go - where - go Nil = Nil - go t@(Tip k x) = if predicate k x then t else Nil - go (Bin p m l r) = bin p m (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'. +-- | /O(n+m)/. Check whether the key sets of two maps are disjoint +-- (i.e. their 'intersection' is empty). -- --- > partitionWithKey (\ k _ -> k > 3) (fromList [(5,"a"), (3,"b")]) == (singleton 5 "a", singleton 3 "b") --- > partitionWithKey (\ k _ -> k < 7) (fromList [(5,"a"), (3,"b")]) == (fromList [(3, "b"), (5, "a")], empty) --- > partitionWithKey (\ k _ -> k > 7) (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3, "b"), (5, "a")]) - -partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a,IntMap a) -partitionWithKey predicate0 t0 = toPair $ go predicate0 t0 - where - go predicate t = - case t of - Bin p m l r -> - let (l1 :*: l2) = go predicate l - (r1 :*: r2) = go predicate r - in bin p m l1 r1 :*: bin p m l2 r2 - Tip k x - | predicate k x -> (t :*: Nil) - | otherwise -> (Nil :*: t) - Nil -> (Nil :*: Nil) - --- | /O(n)/. Map values and collect the 'Just' results. --- --- > let f x = if x == "a" then Just "new a" else Nothing --- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a" - -mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b -mapMaybe f = mapMaybeWithKey (\_ x -> f x) - --- | /O(n)/. Map keys\/values and collect the 'Just' results. --- --- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing --- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" - -mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b -mapMaybeWithKey f (Bin p m l r) - = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r) -mapMaybeWithKey f (Tip k x) = case f k x of - Just y -> Tip k y - Nothing -> Nil -mapMaybeWithKey _ Nil = Nil - --- | /O(n)/. Map values and separate the 'Left' and 'Right' results. --- --- > let f a = if a < "c" then Left a else Right a --- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")]) --- > --- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) - -mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) -mapEither f m - = mapEitherWithKey (\_ x -> f x) m - --- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- > disjoint (fromList [(2,'a')]) (fromList [(1,()), (3,())]) == True +-- > disjoint (fromList [(2,'a')]) (fromList [(1,'a'), (2,'b')]) == False +-- > disjoint (fromList []) (fromList []) == True -- --- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) --- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")]) --- > --- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) - -mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) -mapEitherWithKey f0 t0 = toPair $ go f0 t0 - where - go f (Bin p m l r) = - bin p m l1 r1 :*: bin p m l2 r2 - where - (l1 :*: l2) = go f l - (r1 :*: r2) = go f r - go f (Tip k x) = case f k x of - Left y -> (Tip k y :*: Nil) - Right z -> (Nil :*: Tip k z) - go _ Nil = (Nil :*: Nil) - --- | /O(min(n,W))/. The expression (@'split' k map@) is a pair @(map1,map2)@ --- where all keys in @map1@ are lower than @k@ and all keys in --- @map2@ larger than @k@. Any key equal to @k@ is found in neither @map1@ nor @map2@. +-- > disjoint a b == null (intersection a b) -- --- > split 2 (fromList [(5,"a"), (3,"b")]) == (empty, fromList [(3,"b"), (5,"a")]) --- > split 3 (fromList [(5,"a"), (3,"b")]) == (empty, singleton 5 "a") --- > split 4 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", singleton 5 "a") --- > split 5 (fromList [(5,"a"), (3,"b")]) == (singleton 3 "b", empty) --- > split 6 (fromList [(5,"a"), (3,"b")]) == (fromList [(3,"b"), (5,"a")], empty) - -split :: Key -> IntMap a -> (IntMap a, IntMap a) -split k t = - case t of - Bin _ m l r - | m < 0 -> - if k >= 0 -- handle negative numbers. - then - case go k l of - (lt :*: gt) -> - let !lt' = union r lt - in (lt', gt) - else - case go k r of - (lt :*: gt) -> - let !gt' = union gt l - in (lt, gt') - _ -> case go k t of - (lt :*: gt) -> (lt, gt) +-- @since 0.6.2.1 +disjoint :: IntMap a -> IntMap b -> Bool +disjoint = start where - go k' t'@(Bin p m l r) - | nomatch k' p m = if k' > p then t' :*: Nil else Nil :*: t' - | zero k' m = case go k' l of (lt :*: gt) -> lt :*: union gt r - | otherwise = case go k' r of (lt :*: gt) -> union l lt :*: gt - go k' t'@(Tip ky _) - | k' > ky = (t' :*: Nil) - | k' < ky = (Nil :*: t') - | otherwise = (Nil :*: Nil) - go _ Nil = (Nil :*: Nil) - - -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) + start (IntMap Empty) !_ = True + start !_ (IntMap Empty) = True + start (IntMap (NonEmpty min1 _ root1)) (IntMap (NonEmpty min2 _ root2)) + | min1 < min2 = goL min2 root2 min1 root1 + | min1 > min2 = goL min1 root1 min2 root2 + | otherwise = False + + goL :: Bound L -> Node L x -> Bound L -> Node L y -> Bool + goL !_ !_ !_ Tip = True + goL min1 Tip min2 n2 = goLookupL (boundKey min1) (xor (boundKey min1) min2) n2 + goL min1 (Bin _ _ _ _) _ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = True + goL min1 n1@(Bin max1 _ l1 r1) min2 n2@(Bin max2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> goR max2 r2 max1 (Bin min1 dummyV l1 r1) + | max1 < max2 -> goR max1 (Bin min1 dummyV l1 r1) max2 r2 + | otherwise -> False + EQ | max1 > max2 -> goL min1 l1 min2 l2 && goR max2 r2 max1 r1 + | max1 < max2 -> goL min1 l1 min2 l2 && goR max1 r1 max2 r2 + | otherwise -> False + GT -> goL min1 l1 min2 n2 + + goR :: Bound R -> Node R x -> Bound R -> Node R y -> Bool + goR !_ !_ !_ Tip = True + goR max1 Tip max2 n2 = goLookupR (boundKey max1) (xor (boundKey max1) max2) n2 + goR max1 (Bin _ _ _ _) _ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = True + goR max1 n1@(Bin min1 _ l1 r1) max2 n2@(Bin min2 _ l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> goL min2 l2 min1 (Bin max1 dummyV l1 r1) + | min1 > min2 -> goL min1 (Bin max1 dummyV l1 r1) min2 l2 + | otherwise -> False + EQ | min1 < min2 -> goL min2 l2 min1 l1 && goR max1 r1 max2 r2 + | min1 > min2 -> goL min1 l1 min2 l2 && goR max1 r1 max2 r2 + | otherwise -> False + GT -> goR max1 r1 max2 n2 + + goLookupL !_ !_ Tip = True + goLookupL k !xorCache (Bin max _ l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL k xorCache l + | otherwise -> goLookupR k xorCacheMax r + OutOfBound -> True + Matched -> False + where xorCacheMax = xor k max + + goLookupR !_ !_ Tip = True + goLookupR k !xorCache (Bin min _ l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR k xorCache r + | otherwise -> goLookupL k xorCacheMin l + OutOfBound -> True + Matched -> False + where xorCacheMin = xor k min + + dummyV = error "impossible" -splitLookup :: Key -> IntMap a -> (IntMap a, Maybe a, IntMap a) -splitLookup k t = - case - case t of - Bin _ m l r - | m < 0 -> - if k >= 0 -- handle negative numbers. - then mapLT (union r) (go k l) - else mapGT (`union` l) (go k r) - _ -> go k t - of SplitLookup lt fnd gt -> (lt, fnd, gt) - where - go k' t'@(Bin p m l r) - | nomatch k' p m = - if k' > p - then SplitLookup t' Nothing Nil - else SplitLookup Nil Nothing t' - | zero k' m = mapGT (`union` r) (go k' l) - | otherwise = mapLT (union 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'@. -- @@ -2788,34 +1803,18 @@ splitLookup k t = -- -- > let f a len = len + (length a) -- > foldr f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 -foldr :: (a -> b -> b) -> b -> IntMap a -> b -foldr f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z l) r -- put negative numbers before - | otherwise -> go (go z r) l - _ -> go z t - where - go z' Nil = z' - go z' (Tip _ x) = f x z' - go z' (Bin _ _ l r) = go (go z' r) l {-# INLINE foldr #-} - --- | /O(n)/. A strict version of 'foldr'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldr' :: (a -> b -> b) -> b -> IntMap a -> b -foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z l) r -- put negative numbers before - | otherwise -> go (go z r) l - _ -> go z t +foldr :: (a -> b -> b) -> b -> IntMap a -> b +foldr f z = start where - go !z' Nil = z' - go z' (Tip _ x) = f x z' - go z' (Bin _ _ l r) = go (go z' r) l -{-# INLINE foldr' #-} + start (IntMap Empty) = z + start (IntMap (NonEmpty _ minV root)) = f minV (goL root z) + + goL Tip acc = acc + goL (Bin _ maxV l r) acc = goL l (goR r (f maxV acc)) + + goR Tip acc = acc + goR (Bin _ minV l r) acc = f minV (goL l (goR r acc)) -- | /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'@. @@ -2826,34 +1825,18 @@ foldr' f z = \t -> -- Use lambda t to be inlinable with two arguments only. -- -- > let f len a = len + (length a) -- > foldl f 0 (fromList [(5,"a"), (3,"bbb")]) == 4 +{-# INLINE foldl #-} foldl :: (a -> b -> a) -> a -> IntMap b -> a -foldl f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z r) l -- put negative numbers before - | otherwise -> go (go z l) r - _ -> go z t +foldl f z = start where - go z' Nil = z' - go z' (Tip _ x) = f z' x - go z' (Bin _ _ l r) = go (go z' l) r -{-# INLINE foldl #-} + start (IntMap Empty) = z + start (IntMap (NonEmpty _ minV root)) = goL (f z minV) root --- | /O(n)/. A strict version of 'foldl'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldl' :: (a -> b -> a) -> a -> IntMap b -> a -foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z r) l -- put negative numbers before - | otherwise -> go (go z l) r - _ -> go z t - where - go !z' Nil = z' - go z' (Tip _ x) = f z' x - go z' (Bin _ _ l r) = go (go z' l) r -{-# INLINE foldl' #-} + goL acc Tip = acc + goL acc (Bin _ maxV l r) = f (goR (goL acc l) r) maxV + + goR acc Tip = acc + goR acc (Bin _ minV l r) = goR (goL (f acc minV) l) r -- | /O(n)/. Fold the keys and values in the map using the given right-associative -- binary operator, such that @@ -2865,34 +1848,18 @@ foldl' f z = \t -> -- Use lambda t to be inlinable with two arguments only. -- -- > let f k a result = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" -- > foldrWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (5:a)(3:b)" +{-# INLINE foldrWithKey #-} foldrWithKey :: (Key -> a -> b -> b) -> b -> IntMap a -> b -foldrWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z l) r -- put negative numbers before - | otherwise -> go (go z r) l - _ -> go z t +foldrWithKey f z = start 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 #-} + start (IntMap Empty) = z + start (IntMap (NonEmpty min minV root)) = f (boundKey min) minV (goL root z) --- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b -foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z l) r -- put negative numbers before - | otherwise -> go (go z r) l - _ -> go z t - where - 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' #-} + goL Tip acc = acc + goL (Bin max maxV l r) acc = goL l (goR r (f (boundKey max) maxV acc)) + + goR Tip acc = acc + goR (Bin min minV l r) acc = f (boundKey min) minV (goL l (goR r acc)) -- | /O(n)/. Fold the keys and values in the map using the given left-associative -- binary operator, such that @@ -2904,62 +1871,93 @@ foldrWithKey' f z = \t -> -- Use lambda t to be inlinable with two argument -- -- > let f result k a = result ++ "(" ++ (show k) ++ ":" ++ a ++ ")" -- > foldlWithKey f "Map: " (fromList [(5,"a"), (3,"b")]) == "Map: (3:b)(5:a)" +{-# INLINE foldlWithKey #-} foldlWithKey :: (a -> Key -> b -> a) -> a -> IntMap b -> a -foldlWithKey f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z r) l -- put negative numbers before - | otherwise -> go (go z l) r - _ -> go z t +foldlWithKey f z = start 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 #-} + start (IntMap Empty) = z + start (IntMap (NonEmpty min minV root)) = goL (f z (boundKey min) minV) root --- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is --- evaluated before using the result in the next application. This --- function is strict in the starting value. -foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a -foldlWithKey' f z = \t -> -- Use lambda t to be inlinable with two arguments only. - case t of - Bin _ m l r - | m < 0 -> go (go z r) l -- put negative numbers before - | otherwise -> go (go z l) r - _ -> go z t - where - 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' #-} + goL acc Tip = acc + goL acc (Bin max maxV l r) = f (goR (goL acc l) r) (boundKey max) maxV + + goR acc Tip = acc + goR acc (Bin min minV l r) = goR (goL (f acc (boundKey min) minV) l) r -- | /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 +{-# INLINE foldMapWithKey #-} foldMapWithKey :: Monoid m => (Key -> a -> m) -> IntMap a -> m -foldMapWithKey f = go +foldMapWithKey f = start where - go Nil = mempty - go (Tip kx x) = f kx x - go (Bin _ m l r) - | m < 0 = go r `mappend` go l - | otherwise = go l `mappend` go r -{-# INLINE foldMapWithKey #-} + start (IntMap Empty) = mempty + start (IntMap (NonEmpty min minV root)) = f (boundKey min) minV `mappend` goL root + + goL Tip = mempty + goL (Bin max maxV l r) = goL l `mappend` goR r `mappend` f (boundKey max) maxV + + goR Tip = mempty + goR (Bin min minV l r) = f (boundKey min) minV `mappend` goL l `mappend` goR r + +-- | /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. +{-# INLINE foldr' #-} +foldr' :: (a -> b -> b) -> b -> IntMap a -> b +foldr' f !z = foldrWithKey' (\ _ v !acc -> f v acc) z + +-- | /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. +{-# INLINE foldl' #-} +foldl' :: (a -> b -> a) -> a -> IntMap b -> a +foldl' f !z = foldlWithKey' (\ !acc _ v -> f acc v) z + +-- | /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. +{-# INLINE foldrWithKey' #-} -- Very important for unwrappable accumulators (e.g. Int) +foldrWithKey' :: (Key -> a -> b -> b) -> b -> IntMap a -> b +foldrWithKey' f !z = start + where + f' k v !acc = f k v acc + + start (IntMap Empty) = z + start (IntMap (NonEmpty min minV root)) = f' (boundKey min) minV (goL root z) + + goL Tip !acc = acc + goL (Bin max maxV l r) !acc = goL l (goR r (f' (boundKey max) maxV acc)) + + goR Tip !acc = acc + goR (Bin min minV l r) !acc = f' (boundKey min) minV (goL l (goR r acc)) + +-- | /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. +{-# INLINE foldlWithKey' #-} +foldlWithKey' :: (a -> Key -> b -> a) -> a -> IntMap b -> a +foldlWithKey' f !z = start + where + f' !acc k v = f acc k v + + start (IntMap Empty) = z + start (IntMap (NonEmpty min minV root)) = goL (f' z (boundKey min) minV) root + + goL acc Tip = acc + goL acc (Bin max maxV l r) = f' (goR (goL acc l) r) (boundKey max) maxV + + goR acc Tip = acc + goR acc (Bin min minV l r) = goR (goL (f' acc (boundKey min) minV) l) r -{-------------------------------------------------------------------- - 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 (:) [] @@ -2968,106 +1966,46 @@ elems = foldr (:) [] -- -- > keys (fromList [(5,"a"), (3,"b")]) == [3,5] -- > keys empty == [] - -keys :: IntMap a -> [Key] -keys = foldrWithKey (\k _ ks -> k : ks) [] +keys :: IntMap a -> [Key] +keys = foldrWithKey (\k _ l -> k : l) [] -- | /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 :: IntMap a -> [(Key, a)] assocs = toAscList -- | /O(n*min(n,W))/. The set of all keys of the map. -- -- > keysSet (fromList [(5,"a"), (3,"b")]) == Data.IntSet.fromList [3,5] -- > keysSet empty == Data.IntSet.empty +keysSet :: IntMap a -> Data.IntSet.IntSet +keysSet = Data.IntSet.fromDistinctAscList . keys -keysSet :: IntMap a -> IntSet.IntSet -keysSet Nil = IntSet.Nil -keysSet (Tip kx _) = IntSet.singleton kx -keysSet (Bin p m l r) - | m .&. IntSet.suffixBitMask == 0 = IntSet.Bin p m (keysSet l) (keysSet r) - | otherwise = IntSet.Tip (p .&. IntSet.prefixBitMask) (computeBm (computeBm 0 l) r) - where computeBm !acc (Bin _ _ l' r') = computeBm (computeBm acc l') r' - computeBm acc (Tip kx _) = acc .|. IntSet.bitmapOf kx - computeBm _ Nil = error "Data.IntSet.keysSet: Nil" - --- | /O(n)/. Build a map from a set of keys and a function which for each key --- computes its value. --- --- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] --- > fromSet undefined Data.IntSet.empty == empty - -fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a -fromSet _ IntSet.Nil = Nil -fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r) -fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1) - where - -- This is slightly complicated, as we to convert the dense - -- representation of IntSet into tree representation of IntMap. - -- - -- We are given a nonzero bit mask 'bmask' of 'bits' bits with - -- prefix 'prefix'. We split bmask into halves corresponding - -- to left and right subtree. If they are both nonempty, we - -- create a Bin node, otherwise exactly one of them is nonempty - -- and we construct the IntMap from that half. - buildTree g !prefix !bmask bits = case bits of - 0 -> Tip prefix (g prefix) - _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of - bits2 - | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> - buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2 - | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 -> - buildTree g prefix bmask bits2 - | otherwise -> - Bin prefix bits2 - (buildTree g prefix bmask bits2) - (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2) - -{-------------------------------------------------------------------- - Lists ---------------------------------------------------------------------} -#if __GLASGOW_HASKELL__ >= 708 --- | @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)] +-- | /O(n)/. Convert the map to a list of key\/value pairs. +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) [] +toAscList :: IntMap a -> [(Key, a)] +toAscList = foldrWithKey (\k v l -> (k, v) : l) [] -- | /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) [] +toDescList :: IntMap a -> [(Key, a)] +toDescList = foldlWithKey (\l k v -> (k, v) : l) [] -- List fusion for the list generating functions. -#if __GLASGOW_HASKELL__ +#if USE_REWRITE_RULES -- 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. +-- 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 #-} @@ -3075,7 +2013,7 @@ 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 and toList so that we need to fuse only toAscList. {-# INLINE assocs #-} {-# INLINE toList #-} @@ -3089,370 +2027,225 @@ foldlFB = foldlWithKey {-# 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 #-} +{-# RULES +"IntMap.elems" [~1] forall m . elems m = build (\c n -> foldrFB (\_ x xs -> c x xs) n m) +"IntMap.elems back" [1] foldrFB (\_ x xs -> x : xs) [] = elems +"IntMap.keys" [~1] forall m . keys m = build (\c n -> foldrFB (\k _ xs -> c k xs) n m) +"IntMap.keys back" [1] foldrFB (\k _ xs -> k : xs) [] = keys +"IntMap.toAscList" [~1] forall m . toAscList m = build (\c n -> foldrFB (\k x xs -> c (k, x) xs) n m) +"IntMap.toAscList back" [1] foldrFB (\k x xs -> (k, x) : xs) [] = toAscList +"IntMap.toDescList" [~1] forall m . toDescList m = build (\c n -> foldlFB (\xs k x -> c (k, x) xs) n m) +"IntMap.toDescList back" [1] foldlFB (\xs k x -> (k, x) : xs) [] = toDescList + #-} #endif +-- | A stack used in the in-order building of IntMaps. The utilities here don't +-- force any of the values that are passed, and so require both lazy and strict +-- wrappers. See 'Data.IntMap.Lazy.fromDistinctAscList' and its strict +-- counterpart for examples of how to use these functions. +data BuildStack a = Push {-# UNPACK #-} !(Bound L) a !(Node L a) !(BuildStack a) | StackBase --- | /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")] +pushBuildStack :: Word -> Key -> a -> Node R a -> BuildStack a -> BuildStack a +pushBuildStack !xorCache !k v !r (Push min minV l stk) + | xor k min < xorCache = pushBuildStack xorCache k v (Bin min minV l r) stk +pushBuildStack !_ !k v Tip !stk = Push (Bound k) v Tip stk +pushBuildStack !_ !k v (Bin min minV l r) !stk = Push min minV (Bin (Bound k) v l r) stk -fromList :: [(Key,a)] -> IntMap a -fromList xs - = Foldable.foldl' ins empty xs - where - ins t (k,x) = insert k x t +completeBuildStack :: Bound R -> a -> Node R a -> BuildStack a -> IntMap_ L a +completeBuildStack !max maxV !r StackBase = r2lMap (NonEmpty max maxV r) +completeBuildStack !max maxV !r (Push min minV l stk) = completeBuildStack max maxV (Bin min minV l r) stk --- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- | /O(n)/. Filter all values that satisfy some predicate. -- --- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")] --- > fromListWith (++) [] == empty - -fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWith f xs - = fromListWithKey (\_ x y -> f x y) xs +-- > 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 (const p) --- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. +-- | /O(n)/. Filter all keys\/values that satisfy some predicate. -- --- > 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 +-- > filterWithKey (\k _ -> k > 4) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +{-# INLINE filterWithKey #-} +filterWithKey :: (Key -> a -> Bool) -> IntMap a -> IntMap a +filterWithKey p = filterWithUKey (\k a -> p (box k) a) -fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWithKey f xs - = Foldable.foldl' ins empty xs +-- | /O(n)/. Filter all keys\/values that satisfy some predicate taking unboxed +-- keys. Identical in functionality to 'filterWithKeys'. +filterWithUKey :: (UKey -> a -> Bool) -> IntMap a -> IntMap a +filterWithUKey p = start where - ins t (k,x) = insertWithKey f k x t + start (IntMap Empty) = IntMap Empty + start (IntMap (NonEmpty min minV root)) + | p (boundUKey min) minV = IntMap (NonEmpty min minV (goL root)) + | otherwise = IntMap (goDeleteL root) + + goL Tip = Tip + goL (Bin max maxV l r) + | p (boundUKey max) maxV = Bin max maxV (goL l) (goR r) + | otherwise = binNodeMapL (goL l) (goDeleteR r) + + goR Tip = Tip + goR (Bin min minV l r) + | p (boundUKey min) minV = Bin min minV (goL l) (goR r) + | otherwise = binMapNodeR (goDeleteL l) (goR r) + + goDeleteL Tip = Empty + goDeleteL (Bin max maxV l r) + | p (boundUKey max) maxV = binL (goDeleteL l) (NonEmpty max maxV (goR r)) + | otherwise = binL (goDeleteL l) (goDeleteR r) + + goDeleteR Tip = Empty + goDeleteR (Bin min minV l r) + | p (boundUKey min) minV = binR (NonEmpty min minV (goL l)) (goDeleteR r) + | otherwise = binR (goDeleteL l) (goDeleteR r) --- | /O(n)/. Build a map from a list of key\/value pairs where --- the keys are in ascending order. +-- | /O(n+m)/. The restriction of a map to the keys in a set. -- --- > 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. --- /The precondition (input list is ascending) is not checked./ +-- @ +-- m `restrictKeys` s = 'filterWithKey' (\k _ -> k `'IntSet.member'` s) m +-- @ -- --- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] - -fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f = fromMonoListWithKey Nondistinct (\_ x y -> f x y) -{-# NOINLINE fromAscListWith #-} +-- @since 0.5.8 +restrictKeys :: IntMap a -> Data.IntSet.IntSet -> IntMap a +restrictKeys m s = filterWithKey (\k _ -> Data.IntSet.member k s) m --- | /O(n)/. Build a map from a list of key\/value pairs where --- the keys are in ascending order, with a combining function on equal keys. --- /The precondition (input list is ascending) is not checked./ +-- | Remove all the keys in a given set from a map. -- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")] - -fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey 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. --- /The precondition (input list is strictly ascending) is not checked./ +-- @ +-- m `withoutKeys` s = 'filterWithKey' (\k _ -> k `'IntSet.notMember'` s) m +-- @ -- --- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] - -fromDistinctAscList :: [(Key,a)] -> IntMap a -fromDistinctAscList = fromMonoListWithKey Distinct (\_ x _ -> x) -{-# NOINLINE fromDistinctAscList #-} +-- @since 0.5.8 +withoutKeys :: IntMap a -> Data.IntSet.IntSet -> IntMap a +withoutKeys m s = filterWithKey (\k _ -> Data.IntSet.notMember k s) m --- | /O(n)/. Build a map from a list of key\/value pairs with monotonic keys --- and a combining function. +-- | /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'. -- --- 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. +-- > 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 (const p) -fromMonoListWithKey :: Distinct -> (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromMonoListWithKey distinct f = go +-- | /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")]) +{-# INLINE partitionWithKey #-} +partitionWithKey :: (Key -> a -> Bool) -> IntMap a -> (IntMap a, IntMap a) +partitionWithKey p = partitionWithUKey (\k a -> p (box k) a) + +-- | /O(n)/. Partition the map according to some predicate taking an unboxed +-- key. Identical in functionality to 'partitionWithKey'. +partitionWithUKey :: (UKey -> a -> Bool) -> IntMap a -> (IntMap a, IntMap a) +partitionWithUKey p = start 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 - t1 == t2 = equal t1 t2 - t1 /= t2 = nequal t1 t2 - -equal :: Eq a => IntMap a -> IntMap a -> Bool -equal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 == m2) && (p1 == p2) && (equal l1 l2) && (equal r1 r2) -equal (Tip kx x) (Tip ky y) - = (kx == ky) && (x==y) -equal Nil Nil = True -equal _ _ = False - -nequal :: Eq a => IntMap a -> IntMap a -> Bool -nequal (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 /= m2) || (p1 /= p2) || (nequal l1 l2) || (nequal r1 r2) -nequal (Tip kx x) (Tip ky y) - = (kx /= ky) || (x/=y) -nequal Nil Nil = False -nequal _ _ = True - -#if MIN_VERSION_base(4,9,0) --- | @since 0.5.9 -instance Eq1 IntMap where - liftEq eq (Bin p1 m1 l1 r1) (Bin p2 m2 l2 r2) - = (m1 == m2) && (p1 == p2) && (liftEq eq l1 l2) && (liftEq eq r1 r2) - liftEq eq (Tip kx x) (Tip ky y) - = (kx == ky) && (eq x y) - liftEq _eq Nil Nil = True - liftEq _eq _ _ = False -#endif - -{-------------------------------------------------------------------- - Ord ---------------------------------------------------------------------} - -instance Ord a => Ord (IntMap a) where - compare m1 m2 = compare (toList m1) (toList m2) - -#if MIN_VERSION_base(4,9,0) --- | @since 0.5.9 -instance Ord1 IntMap where - liftCompare cmp m n = - liftCompare (liftCompare cmp) (toList m) (toList n) -#endif - -{-------------------------------------------------------------------- - Functor ---------------------------------------------------------------------} - -instance Functor IntMap where - fmap = map - -#ifdef __GLASGOW_HASKELL__ - a <$ Bin p m l r = Bin p m (a <$ l) (a <$ r) - a <$ Tip k _ = Tip k a - _ <$ Nil = Nil -#endif - -{-------------------------------------------------------------------- - Show ---------------------------------------------------------------------} - -instance Show a => Show (IntMap a) where - showsPrec d m = showParen (d > 10) $ - showString "fromList " . shows (toList m) - -#if MIN_VERSION_base(4,9,0) --- | @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 -#endif - -{-------------------------------------------------------------------- - Read ---------------------------------------------------------------------} -instance (Read e) => Read (IntMap e) where -#ifdef __GLASGOW_HASKELL__ - readPrec = parens $ prec 10 $ do - Ident "fromList" <- lexP - xs <- readPrec - return (fromList xs) - - readListPrec = readListPrecDefault -#else - readsPrec p = readParen (p > 10) $ \ r -> do - ("fromList",s) <- lex r - (xs,t) <- reads s - return (fromList xs,t) -#endif - -#if MIN_VERSION_base(4,9,0) --- | @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 -#endif + start (IntMap Empty) = (IntMap Empty, IntMap Empty) + start (IntMap (NonEmpty min minV root)) + | p (boundUKey min) minV = let t :*: f = goTrueL root + in (IntMap (NonEmpty min minV t), IntMap f) + | otherwise = let t :*: f = goFalseL root + in (IntMap t, IntMap (NonEmpty min minV f)) + + goTrueL Tip = Tip :*: Empty + goTrueL (Bin max maxV l r) + | p (boundUKey max) maxV = let tl :*: fl = goTrueL l + tr :*: fr = goTrueR r + in Bin max maxV tl tr :*: binL fl fr + | otherwise = let tl :*: fl = goTrueL l + tr :*: fr = goFalseR r + in binNodeMapL tl tr :*: binL fl (NonEmpty max maxV fr) + + goTrueR Tip = Tip :*: Empty + goTrueR (Bin min minV l r) + | p (boundUKey min) minV = let tl :*: fl = goTrueL l + tr :*: fr = goTrueR r + in Bin min minV tl tr :*: binR fl fr + | otherwise = let tl :*: fl = goFalseL l + tr :*: fr = goTrueR r + in binMapNodeR tl tr :*: binR (NonEmpty min minV fl) fr + + goFalseL Tip = Empty :*: Tip + goFalseL (Bin max maxV l r) + | p (boundUKey max) maxV = + let tl :*: fl = goFalseL l + tr :*: fr = goTrueR r + in binL tl (NonEmpty max maxV tr) :*: binNodeMapL fl fr + | otherwise = let tl :*: fl = goFalseL l + tr :*: fr = goFalseR r + in binL tl tr :*: Bin max maxV fl fr + + goFalseR Tip = Empty :*: Tip + goFalseR (Bin min minV l r) + | p (boundUKey min) minV = + let tl :*: fl = goTrueL l + tr :*: fr = goFalseR r + in binR (NonEmpty min minV tl) tr :*: binMapNodeR fl fr + | otherwise = let tl :*: fl = goFalseL l + tr :*: fr = goFalseR r + in binR tl tr :*: Bin min minV fl fr -{-------------------------------------------------------------------- - Typeable ---------------------------------------------------------------------} - -INSTANCE_TYPEABLE1(IntMap) +-- | /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 m = case splitLookup k m of + (lt, _, gt) -> (lt, gt) -{-------------------------------------------------------------------- - Helpers ---------------------------------------------------------------------} -{-------------------------------------------------------------------- - Link ---------------------------------------------------------------------} -link :: Prefix -> IntMap a -> Prefix -> IntMap a -> IntMap a -link p1 t1 p2 t2 = linkWithMask (branchMask p1 p2) p1 t1 {-p2-} t2 -{-# INLINE link #-} - --- `linkWithMask` is useful when the `branchMask` has already been computed -linkWithMask :: Mask -> Prefix -> IntMap a -> IntMap a -> IntMap a -linkWithMask m p1 t1 {-p2-} t2 - | zero p1 m = Bin p m t1 t2 - | otherwise = Bin p m t2 t1 +-- | /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 = start where - p = mask p1 m -{-# INLINE linkWithMask #-} - -{-------------------------------------------------------------------- - @bin@ assures that we never have empty trees within a tree. ---------------------------------------------------------------------} -bin :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a -bin _ _ l Nil = l -bin _ _ Nil r = r -bin p m l r = Bin p m l r -{-# INLINE bin #-} - --- binCheckLeft only checks that the left subtree is non-empty -binCheckLeft :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a -binCheckLeft _ _ Nil r = r -binCheckLeft p m l r = Bin p m l r -{-# INLINE binCheckLeft #-} - --- binCheckRight only checks that the right subtree is non-empty -binCheckRight :: Prefix -> Mask -> IntMap a -> IntMap a -> IntMap a -binCheckRight _ _ l Nil = l -binCheckRight p m l r = Bin p m l r -{-# INLINE binCheckRight #-} - -{-------------------------------------------------------------------- - Endian independent bit twiddling ---------------------------------------------------------------------} - --- | Should this key follow the left subtree of a 'Bin' with switching --- bit @m@? N.B., the answer is only valid when @match i p m@ is true. -zero :: Key -> Mask -> Bool -zero i m - = (natFromInt i) .&. (natFromInt m) == 0 -{-# INLINE zero #-} - -nomatch,match :: Key -> Prefix -> Mask -> Bool - --- | Does the key @i@ differ from the prefix @p@ before getting to --- the switching bit @m@? -nomatch i p m - = (mask i m) /= p -{-# INLINE nomatch #-} - --- | Does the key @i@ match the prefix @p@ (up to but not including --- bit @m@)? -match i p m - = (mask i m) == p -{-# INLINE match #-} - - --- | The prefix of key @i@ up to (but not including) the switching --- bit @m@. -mask :: Key -> Mask -> Prefix -mask i m - = maskW (natFromInt i) (natFromInt m) -{-# INLINE mask #-} - - -{-------------------------------------------------------------------- - Big endian operations ---------------------------------------------------------------------} - --- | The prefix of key @i@ up to (but not including) the switching --- bit @m@. -maskW :: Nat -> Nat -> Prefix -maskW i m - = intFromNat (i .&. ((-m) `xor` m)) -{-# INLINE maskW #-} - --- | Does the left switching bit specify a shorter prefix? -shorter :: Mask -> Mask -> Bool -shorter m1 m2 - = (natFromInt m1) > (natFromInt m2) -{-# INLINE shorter #-} - --- | The first switching bit where the two prefixes disagree. -branchMask :: Prefix -> Prefix -> Mask -branchMask p1 p2 - = intFromNat (highestBitMask (natFromInt p1 `xor` natFromInt p2)) -{-# INLINE branchMask #-} - -{-------------------------------------------------------------------- - 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. + start (IntMap Empty) = (IntMap Empty, Nothing, IntMap Empty) + start m@(IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> case root of + Tip -> (m, Nothing, IntMap Empty) + Bin max maxV l r -> case compareMaxBound k max of + InBound -> let (NE glb glbV lt, eq, NE lub lubV gt) = go (xor k min) min minV (xor k max) max maxV l r + in (IntMap (r2lMap (NonEmpty glb glbV lt)), eq, IntMap (NonEmpty lub lubV gt)) + OutOfBound -> (m, Nothing, IntMap Empty) + Matched -> let NE max' maxV' root' = deleteMaxR min minV l r + in (IntMap (r2lMap (NonEmpty max' maxV' root')), Just maxV, IntMap Empty) + OutOfBound -> (IntMap Empty, Nothing, m) + Matched -> (IntMap Empty, Just minV, IntMap (nodeToMapL root)) + + go xorCacheMin min minV xorCacheMax max maxV l r + | xorCacheMin < xorCacheMax = case l of + Tip -> (NE (minToMax min) minV Tip, Nothing, r2lNE (NE max maxV r)) + Bin maxI maxVI lI rI -> case compareMaxBound k maxI of + InBound -> let (lt, eq, NE minI minVI gt) = go xorCacheMin min minV (xor k maxI) maxI maxVI lI rI + in (lt, eq, NE minI minVI (Bin max maxV gt r)) + OutOfBound -> (l2rNE (NE min minV l), Nothing, r2lNE (NE max maxV r)) + Matched -> (deleteMaxR min minV lI rI, Just maxVI, r2lNE (NE max maxV r)) + | otherwise = case r of + Tip -> (l2rNE (NE min minV l), Nothing, NE (maxToMin max) maxV Tip) + Bin minI minVI lI rI -> case compareMinBound k minI of + InBound -> let (NE maxI maxVI lt, eq, gt) = go (xor k minI) minI minVI xorCacheMax max maxV lI rI + in (NE maxI maxVI (Bin min minV l lt), eq, gt) + OutOfBound -> (l2rNE (NE min minV l), Nothing, r2lNE (NE max maxV r)) + Matched -> (l2rNE (NE min minV l), Just minVI, deleteMinL max maxV lI rI) + +-- | /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 @@ -3469,81 +2262,482 @@ branchMask p1 p2 -- Note that the current implementation does not return more than two submaps, -- but you should not depend on this behaviour because it can change in the -- future without notice. -splitRoot :: IntMap a -> [IntMap a] -splitRoot orig = - case orig of - Nil -> [] - x@(Tip _ _) -> [x] - Bin _ m l r | m < 0 -> [r, l] - | otherwise -> [l, r] {-# INLINE splitRoot #-} +splitRoot :: IntMap a -> [IntMap a] +splitRoot (IntMap Empty) = [] +splitRoot m@(IntMap (NonEmpty _ _ Tip)) = [m] +splitRoot (IntMap (NonEmpty min minV (Bin max maxV l r))) = [IntMap (NonEmpty min minV l), IntMap (r2lMap (NonEmpty max maxV r))] + +-- | /O(n+m)/. Is this a submap? +-- Defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +isSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool +isSubmapOf = isSubmapOfBy (==) + +{- | /O(n+m)/. + The expression (@'isSubmapOfBy' f m1 m2@) returns 'True' if + all keys in @m1@ are in @m2@, and when @f@ returns 'True' when + applied to their respective values. For example, the following + expressions are all 'True': + + > isSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1),(2,2)]) + + But the following are all 'False': + + > isSubmapOfBy (==) (fromList [(1,2)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (<) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) + > isSubmapOfBy (==) (fromList [(1,1),(2,2)]) (fromList [(1,1)]) +-} +isSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool +isSubmapOfBy p = start + where + -- 'isSubmapOfBy' follows the pattern of 'intersection' with a few exceptions: + -- + -- * We compare bounds (@max1@ and @max2@ in 'goL' and @min1@ and @min2@ in + -- 'goR') before we compare MSBs, since if a subtree of map 1 isn't + -- entirely contained withing map 2, we can immediately return 'False'. + -- * The "disjoint bounds" case is dropped, as it is already caught by + -- comparing minima (in 'goR') or maxima (in 'goL') in the next step. + -- * Once we know that map 1 has a smaller range than map 2 (that is, + -- @min1 > min2 && max1 < max2@, we know map 1 can't have an earlier + -- splitting bit, so we can use the simpler 'ltMSB' instead of + -- 'compareMSB'. + -- * The preconditions of @goL2@, @goR2@, and such in 'intersection' + -- imply that map 1 isn't a submap of map 2, so we omit those and just + -- return 'False'. + start (IntMap Empty) !_ = True + start !_ (IntMap Empty) = False + start (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = False + | min1 > min2 = goL minV1 min1 root1 min2 root2 + | otherwise = p minV1 minV2 && goLFused min1 root1 root2 + + goL minV1 min1 Tip min2 n2 = goLookupL (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL _ _ _ _ Tip = False + goL minV1 min1 n1@(Bin max1 maxV1 l1 r1) min2 (Bin max2 maxV2 l2 r2) + | max1 > max2 = False + | max1 < max2 = case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL minV1 min1 n1 min2 l2 -- LT + | otherwise -> goR maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + False -> goL minV1 min1 l1 min2 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | otherwise = p maxV1 maxV2 && case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True -> goRFused max1 (Bin min1 minV1 l1 r1) r2 -- LT + False -> goL minV1 min1 l1 min2 l2 && goRFused max1 r1 r2 -- EQ + + goLFused _ Tip _ = True + goLFused _ _ Tip = False + goLFused min n1@(Bin max1 maxV1 l1 r1) (Bin max2 maxV2 l2 r2) + | max1 > max2 = False + | max1 < max2 = case xorBounds min max1 `ltMSB` xorBounds min max2 of + True -> goLFused min n1 l2 + False -> goLFused min l1 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | otherwise = p maxV1 maxV2 && goLFused min l1 l2 && goRFused max1 r1 r2 + + goR maxV1 max1 Tip max2 n2 = goLookupR (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR _ _ _ _ Tip = False + goR maxV1 max1 n1@(Bin min1 minV1 l1 r1) max2 (Bin min2 minV2 l2 r2) + | min1 < min2 = False + | min1 > min2 = case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR maxV1 max1 n1 max2 r2 -- LT + | otherwise -> goL minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + False -> goL minV1 min1 l1 min2 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | otherwise = p minV1 minV2 && case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True -> goLFused min1 (Bin max1 maxV1 l1 r1) l2 -- LT + False -> goLFused min1 l1 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + + goRFused _ Tip _ = True + goRFused _ _ Tip = False + goRFused max n1@(Bin min1 minV1 l1 r1) (Bin min2 minV2 l2 r2) + | min1 < min2 = False + | min1 > min2 = case xorBounds min1 max `ltMSB` xorBounds min2 max of + True -> goRFused max n1 r2 + False -> goL minV1 min1 l1 min2 l2 && goRFused max r1 r2 -- EQ + | otherwise = p minV1 minV2 && goLFused min1 l1 l2 && goRFused max r1 r2 + + goLookupL _ _ !_ Tip = False + goLookupL k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL k v xorCache l + | otherwise -> goLookupR k v xorCacheMax r + OutOfBound -> False + Matched -> p v maxV + where xorCacheMax = xor k max + + goLookupR _ _ !_ Tip = False + goLookupR k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR k v xorCache r + | otherwise -> goLookupL k v xorCacheMin l + OutOfBound -> False + Matched -> p v minV + where xorCacheMin = xor k min +-- | /O(n+m)/. Is this a proper submap? (ie. a submap but not equal). +-- Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' (==)@). +isProperSubmapOf :: Eq a => IntMap a -> IntMap a -> Bool +isProperSubmapOf = isProperSubmapOfBy (==) -{-------------------------------------------------------------------- - Debugging ---------------------------------------------------------------------} +{- | /O(n+m)/. 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': --- | /O(n)/. Show the tree that implements the map. The tree is shown --- in a compressed, hanging format. -showTree :: Show a => IntMap a -> String -showTree s - = showTreeWith True False s +> isProperSubmapOfBy (==) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) +> isProperSubmapOfBy (<=) (fromList [(1,1)]) (fromList [(1,1),(2,2)]) +But the following are all 'False': -{- | /O(n)/. The expression (@'showTreeWith' hang wide map@) shows - the tree that implements the map. If @hang@ is - 'True', a /hanging/ tree is shown otherwise a rotated tree is shown. If - @wide@ is 'True', an extra wide version is shown. +> 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)]) -} -showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String -showTreeWith hang wide t - | hang = (showsTreeHang wide [] t) "" - | otherwise = (showsTree wide [] [] t) "" - -showsTree :: Show a => Bool -> [String] -> [String] -> IntMap a -> ShowS -showsTree wide lbars rbars t = case t of - Bin p m l r -> - showsTree wide (withBar rbars) (withEmpty rbars) r . - showWide wide rbars . - showsBars lbars . showString (showBin p m) . showString "\n" . - showWide wide lbars . - showsTree wide (withEmpty lbars) (withBar lbars) l - Tip k x -> - showsBars lbars . - showString " " . shows k . showString ":=" . shows x . showString "\n" - Nil -> showsBars lbars . showString "|\n" - -showsTreeHang :: Show a => Bool -> [String] -> IntMap a -> ShowS -showsTreeHang wide bars t = case t of - Bin p m l r -> - showsBars bars . showString (showBin p m) . showString "\n" . - showWide wide bars . - showsTreeHang wide (withBar bars) l . - showWide wide bars . - showsTreeHang wide (withEmpty bars) r - Tip k x -> - showsBars bars . - showString " " . shows k . showString ":=" . shows x . showString "\n" - Nil -> showsBars bars . showString "|\n" - -showBin :: Prefix -> Mask -> String -showBin _ _ - = "*" -- ++ show (p,m) - -showWide :: Bool -> [String] -> String -> String -showWide wide bars - | wide = showString (concat (reverse bars)) . showString "|\n" - | otherwise = id - -showsBars :: [String] -> ShowS -showsBars bars - = case bars of - [] -> id - _ -> showString (concat (reverse (tail bars))) . showString node - -node :: String -node = "+--" - -withBar, withEmpty :: [String] -> [String] -withBar bars = "| ":bars -withEmpty bars = " ":bars +isProperSubmapOfBy :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Bool +isProperSubmapOfBy p m1 m2 = submapCmp p m1 m2 == LT + +submapCmp :: (a -> b -> Bool) -> IntMap a -> IntMap b -> Ordering +submapCmp p = start + where + start (IntMap Empty) (IntMap Empty) = EQ + start (IntMap Empty) !_ = LT + start !_ (IntMap Empty) = GT + start (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = GT + | min1 > min2 = fromBool $ goL minV1 min1 root1 min2 root2 + | p minV1 minV2 = goLFused min1 root1 root2 + | otherwise = GT + + goL minV1 min1 Tip min2 n2 = goLookupL (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL _ _ _ _ Tip = False + goL minV1 min1 n1@(Bin max1 maxV1 l1 r1) min2 (Bin max2 maxV2 l2 r2) + | max1 > max2 = False + | max1 < max2 = case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL minV1 min1 n1 min2 l2 -- LT + | otherwise -> goR maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + False -> goL minV1 min1 l1 min2 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | otherwise = p maxV1 maxV2 && case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True -> goRFusedBool max1 (Bin min1 minV1 l1 r1) r2 -- LT + False -> goL minV1 min1 l1 min2 l2 && goRFusedBool max1 r1 r2 -- EQ + + goLFused _ Tip Tip = EQ + goLFused _ Tip _ = LT + goLFused _ _ Tip = GT + goLFused min n1@(Bin max1 maxV1 l1 r1) (Bin max2 maxV2 l2 r2) + | max1 > max2 = GT + | max1 < max2 = fromBool $ case xorBounds min max1 `ltMSB` xorBounds min max2 of + True -> goLFusedBool min n1 l2 + False -> goLFusedBool min l1 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | p maxV1 maxV2 = goLFused min l1 l2 `combine` goRFused max1 r1 r2 + | otherwise = GT + + goLFusedBool _ Tip _ = True + goLFusedBool _ _ Tip = False + goLFusedBool min n1@(Bin max1 maxV1 l1 r1) (Bin max2 maxV2 l2 r2) + | max1 > max2 = False + | max1 < max2 = case xorBounds min max1 `ltMSB` xorBounds min max2 of + True -> goLFusedBool min n1 l2 + False -> goLFusedBool min l1 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | otherwise = p maxV1 maxV2 && goLFusedBool min l1 l2 && goRFusedBool max1 r1 r2 + + goR maxV1 max1 Tip max2 n2 = goLookupR (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR _ _ _ _ Tip = False + goR maxV1 max1 n1@(Bin min1 minV1 l1 r1) max2 (Bin min2 minV2 l2 r2) + | min1 < min2 = False + | min1 > min2 = case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR maxV1 max1 n1 max2 r2 -- LT + | otherwise -> goL minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + False -> goL minV1 min1 l1 min2 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + | otherwise = p minV1 minV2 && case xorBounds min1 max1 `ltMSB` xorBounds min2 max2 of + True -> goLFusedBool min1 (Bin max1 maxV1 l1 r1) l2 -- LT + False -> goLFusedBool min1 l1 l2 && goR maxV1 max1 r1 max2 r2 -- EQ + + goRFused _ Tip Tip = EQ + goRFused _ Tip _ = LT + goRFused _ _ Tip = GT + goRFused max n1@(Bin min1 minV1 l1 r1) (Bin min2 minV2 l2 r2) + | min1 < min2 = GT + | min1 > min2 = fromBool $ case xorBounds min1 max `ltMSB` xorBounds min2 max of + True -> goRFusedBool max n1 r2 + False -> goL minV1 min1 l1 min2 l2 && goRFusedBool max r1 r2 -- EQ + | p minV1 minV2 = goLFused min1 l1 l2 `combine` goRFused max r1 r2 + | otherwise = GT + + goRFusedBool _ Tip _ = True + goRFusedBool _ _ Tip = False + goRFusedBool max n1@(Bin min1 minV1 l1 r1) (Bin min2 minV2 l2 r2) + | min1 < min2 = False + | min1 > min2 = case xorBounds min1 max `ltMSB` xorBounds min2 max of + True -> goRFusedBool max n1 r2 + False -> goL minV1 min1 l1 min2 l2 && goRFusedBool max r1 r2 -- EQ + | otherwise = p minV1 minV2 && goLFusedBool min1 l1 l2 && goRFusedBool max r1 r2 + + goLookupL _ _ !_ Tip = False + goLookupL k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL k v xorCache l + | otherwise -> goLookupR k v xorCacheMax r + OutOfBound -> False + Matched -> p v maxV + where xorCacheMax = xor k max + + goLookupR _ _ !_ Tip = False + goLookupR k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR k v xorCache r + | otherwise -> goLookupL k v xorCacheMin l + OutOfBound -> False + Matched -> p v minV + where xorCacheMin = xor k min + + fromBool True = LT + fromBool False = GT + + combine GT _ = GT + combine _ GT = GT + combine EQ EQ = EQ + combine _ _ = LT + + +-- | /O(1)/. The minimal key of the map. Returns 'Nothing' if the map is empty. +lookupMin :: IntMap a -> Maybe (Key, a) +lookupMin (IntMap Empty) = Nothing +lookupMin (IntMap (NonEmpty min minV _)) = Just (boundKey min, minV) + +-- | /O(1)/. The maximal key of the map. Returns 'Nothing' if the map is empty. +lookupMax :: IntMap a -> Maybe (Key, a) +lookupMax (IntMap Empty) = Nothing +lookupMax (IntMap (NonEmpty min minV root)) = case root of + Tip -> Just (boundKey min, minV) + Bin max maxV _ _ -> Just (boundKey max, maxV) + +-- | /O(1)/. The minimal key of the map. +findMin :: IntMap a -> (Key, a) +findMin (IntMap Empty) = error "findMin: empty map has no minimal element" +findMin (IntMap (NonEmpty min minV _)) = (boundKey min, minV) + +-- | /O(1)/. The maximal key of the map. +findMax :: IntMap a -> (Key, a) +findMax (IntMap Empty) = error "findMin: empty map has no minimal element" +findMax (IntMap (NonEmpty min minV root)) = case root of + Tip -> (boundKey min, minV) + Bin max maxV _ _ -> (boundKey max, maxV) + +-- | /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 (IntMap Empty) = IntMap Empty +deleteMin (IntMap (NonEmpty _ _ root)) = IntMap (nodeToMapL root) + +-- | /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 (IntMap Empty) = IntMap Empty +deleteMax (IntMap (NonEmpty _ _ Tip)) = IntMap Empty +deleteMax (IntMap (NonEmpty min minV (Bin _ _ l r))) = IntMap (NonEmpty min minV (extractBinL l r)) + +-- | /O(min(n,W))/. Delete and find the minimal element. +deleteFindMin :: IntMap a -> ((Key, a), IntMap a) +deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") + . minViewWithKey + +-- | /O(min(n,W))/. Delete and find the maximal element. +deleteFindMax :: IntMap a -> ((Key, a), IntMap a) +deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") + . maxViewWithKey + +-- | /O(min(n,W))/. 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 (IntMap Empty) = Nothing +minView (IntMap (NonEmpty _ minV root)) = Just (minV, IntMap (nodeToMapL root)) + +-- | /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 (IntMap Empty) = Nothing +maxView (IntMap (NonEmpty _ minV Tip)) = Just (minV, IntMap Empty) +maxView (IntMap (NonEmpty min minV (Bin _ maxV l r))) = Just (maxV, IntMap (NonEmpty min minV (extractBinL l r))) + +-- | /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 (IntMap Empty) = Nothing +minViewWithKey (IntMap (NonEmpty min minV root)) = Just ((boundKey min, minV), IntMap (nodeToMapL root)) + +-- | /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 (IntMap Empty) = Nothing +maxViewWithKey (IntMap (NonEmpty min minV Tip)) = Just ((boundKey min, minV), IntMap Empty) +maxViewWithKey (IntMap (NonEmpty min minV (Bin max maxV l r))) = Just ((boundKey max, maxV), IntMap (NonEmpty min minV (extractBinL l r))) + +-- | /O(1)/. Returns whether the most significant bit (MSB) of its first +-- argument is less significant than the most significant bit of its second +-- argument. + +-- This works by measuring whether x is in between 0 and y but closer to 0 (in the XOR metric). +{-# INLINE ltMSB #-} +ltMSB :: Word -> Word -> Bool +ltMSB x y = x < y && x < Data.Bits.xor x y + +-- | /O(1)/. Compares the significance of the most significan set bits (MSBs) +-- of two words. Commonly used with `xorBounds` to tell which of two nodes +-- spans a larger space of keys in merge operations. +-- See 'ltMSB' for why this works +{-# INLINE compareMSB #-} +compareMSB :: Word -> Word -> Ordering +compareMSB x y = case x < y of + True | x < Data.Bits.xor x y -> LT + False | y < Data.Bits.xor x y -> GT + _ -> EQ + +{-# INLINE binL #-} +binL :: IntMap_ L a -> IntMap_ R a -> IntMap_ L a +binL Empty r = r2lMap r +binL l Empty = l +binL (NonEmpty min minV l) (NonEmpty max maxV r) = NonEmpty min minV (Bin max maxV l r) + +{-# INLINE binR #-} +binR :: IntMap_ L a -> IntMap_ R a -> IntMap_ R a +binR Empty r = r +binR l Empty = l2rMap l +binR (NonEmpty min minV l) (NonEmpty max maxV r) = NonEmpty max maxV (Bin min minV l r) + +{-# INLINE binNodeMapL #-} +binNodeMapL :: Node L v -> IntMap_ R v -> Node L v +binNodeMapL l Empty = l +binNodeMapL l (NonEmpty max maxV r) = Bin max maxV l r + +{-# INLINE binMapNodeR #-} +binMapNodeR :: IntMap_ L v -> Node R v -> Node R v +binMapNodeR Empty r = r +binMapNodeR (NonEmpty min minV l) r = Bin min minV l r + +{-# INLINE minToMax #-} +minToMax :: Bound L -> Bound R +minToMax = Bound . boundKey + +{-# INLINE maxToMin #-} +maxToMin :: Bound R -> Bound L +maxToMin = Bound . boundKey + +{-# INLINE l2rMap #-} +l2rMap :: IntMap_ L a -> IntMap_ R a +l2rMap Empty = Empty +l2rMap (NonEmpty min minV Tip) = NonEmpty (minToMax min) minV Tip +l2rMap (NonEmpty min minV (Bin max maxV l r)) = NonEmpty max maxV (Bin min minV l r) + +{-# INLINE r2lMap #-} +r2lMap :: IntMap_ R a -> IntMap_ L a +r2lMap Empty = Empty +r2lMap (NonEmpty max maxV Tip) = NonEmpty (maxToMin max) maxV Tip +r2lMap (NonEmpty max maxV (Bin min minV l r)) = NonEmpty min minV (Bin max maxV l r) + +{-# INLINE l2rNE #-} +l2rNE :: NonEmptyIntMap_ L a -> NonEmptyIntMap_ R a +l2rNE (NE min minV Tip) = NE (minToMax min) minV Tip +l2rNE (NE min minV (Bin max maxV l r)) = NE max maxV (Bin min minV l r) + +{-# INLINE r2lNE #-} +r2lNE :: NonEmptyIntMap_ R a -> NonEmptyIntMap_ L a +r2lNE (NE max maxV Tip) = NE (maxToMin max) maxV Tip +r2lNE (NE max maxV (Bin min minV l r)) = NE min minV (Bin max maxV l r) + +-- | Insert a key/value pair to a left node where the key is smaller than +-- any present in that node. Requires the XOR of the inserted key and the +-- key immediately prior to it (the minimum bound of the node). +insertMinL :: Word -> Bound L -> a -> Node L a -> Node L a +insertMinL !_ !min minV Tip = Bin (minToMax min) minV Tip Tip +insertMinL !xorCache !min minV (Bin max maxV l r) + -- Although the new minimum is not directly passed into 'insertMinL', + -- it is captured in the 'xorCache'. We use standard navigation to + -- determine whether 'min' should belong in the left or right branch. + -- Since 'min' is, by assumption, smaller than any key in the tree, + -- if 'min' is assigned to the right branch than the entire subtree + -- must fit in the right branch. Otherwise, we need to continue recursing. + | xor (boundKey min) max < xorCache = Bin max maxV Tip (Bin min minV l r) + | otherwise = Bin max maxV (insertMinL xorCache min minV l) r + +-- | Insert a key/value pair to a right node where the key is greater than +-- any present in that node. Requires the XOR of the inserted key and the +-- key immediately following it (the maximum bound of the node). +insertMaxR :: Word -> Bound R -> a -> Node R a -> Node R a +insertMaxR !_ !max maxV Tip = Bin (maxToMin max) maxV Tip Tip +insertMaxR !xorCache !max maxV (Bin min minV l r) + | xor (boundKey max) min < xorCache = Bin min minV (Bin max maxV l r) Tip + | otherwise = Bin min minV l (insertMaxR xorCache max maxV r) + +-- | Rearrange an unpacked (non-empty) left node into a non-empty map, for use +-- when the minimum bound of the node has been deleted. See 'NonEmptyIntMap_' +-- for the reasoning behind this. +deleteMinL :: Bound R -> a -> Node L a -> Node R a -> NonEmptyIntMap_ L a +deleteMinL !max maxV l r = case l of -- force l-then-r match order + Tip -> case r of + Tip -> NE (maxToMin max) maxV Tip + Bin min minV innerL innerR -> NE min minV (Bin max maxV innerL innerR) + Bin innerMax innerMaxV innerL innerR -> + let NE min minV inner = deleteMinL innerMax innerMaxV innerL innerR + in NE min minV (Bin max maxV inner r) + +-- | Rearrange an unpacked (non-empty) right node into a non-empty map, for use +-- when the maximum bound of the node has been deleted. See 'NonEmptyIntMap_' +-- for the reasoning behind this. +deleteMaxR :: Bound L -> a -> Node L a -> Node R a -> NonEmptyIntMap_ R a +deleteMaxR !min minV l r = case r of -- force r-then-l match order + Tip -> case l of + Tip -> NE (minToMax min) minV Tip + Bin max maxV innerL innerR -> NE max maxV (Bin min minV innerL innerR) + Bin innerMin innerMinV innerL innerR -> + let NE max maxV inner = deleteMaxR innerMin innerMinV innerL innerR + in NE max maxV (Bin min minV l inner) + +-- | Combine two nodes that would be on different branches into into a new left +-- node. This is not cheap: since the 'Bin' constructor needs an new bound in +-- addition to two children, that bound must be extracted from a child. The +-- primary use case of this and 'extractBinR' is readjusting a node after the +-- bound that it stores has been deleted, so only the two children remain. +-- +-- 'NonEmptyIntMap_' has more details about this use case. +extractBinL :: Node L a -> Node R a -> Node L a +extractBinL l Tip = l +extractBinL l (Bin min minV innerL innerR) = + let NE max maxV r = deleteMaxR min minV innerL innerR + in Bin max maxV l r + +-- | Combine two nodes that would be on different branches into into a new +-- right node. This is not cheap: since the 'Bin' constructor needs an new +-- bound in addition to two children, that bound must be extracted from a +-- child. The primary use case of this and 'extractBinL' is readjusting a node +-- after the bound that it stores has been deleted, so only the two children +-- remain. +-- +-- 'NonEmptyIntMap_' has more details about this use case. +extractBinR :: Node L a -> Node R a -> Node R a +extractBinR Tip r = r +extractBinR (Bin max maxV innerL innerR) r = + let NE min minV l = deleteMinL max maxV innerL innerR + in Bin min minV l r + +-- | Convert a left 'Node' into an 'IntMap_' for use when the external minimum +-- is no longer available. See 'NonEmptyIntMap_' for more details. +nodeToMapL :: Node L a -> IntMap_ L a +nodeToMapL Tip = Empty +nodeToMapL (Bin max maxV innerL innerR) = + let NE min minV l = deleteMinL max maxV innerL innerR + in NonEmpty min minV l + +-- | Convert a right 'Node' into an 'IntMap_' for use when the external maximum +-- is no longer available. See 'NonEmptyIntMap_' for more details. +nodeToMapR :: Node R a -> IntMap_ R a +nodeToMapR Tip = Empty +nodeToMapR (Bin min minV innerL innerR) = + let NE max maxV r = deleteMaxR min minV innerL innerR + in NonEmpty max maxV r diff --git a/containers/src/Data/IntMap/Internal/Debug.hs b/containers/src/Data/IntMap/Internal/Debug.hs index a30dc6968..c908e4767 100644 --- a/containers/src/Data/IntMap/Internal/Debug.hs +++ b/containers/src/Data/IntMap/Internal/Debug.hs @@ -1,6 +1,131 @@ +{-# LANGUAGE CPP, BangPatterns #-} + module Data.IntMap.Internal.Debug ( showTree , showTreeWith + , valid + , validWith ) where +import Numeric (showIntAtBase) +import Data.Char (intToDigit) +#if !MIN_VERSION_base(4,8,0) +import Data.Word (Word) +#endif + +import Prelude hiding (min, max) import Data.IntMap.Internal + +-- | Show the tree that implements the map. The tree is shown in in-order, +-- ASCII format. +showTree :: Show a => IntMap a -> String +showTree = showTreeWith True False where + +-- | The expression @'showTreeWith' inorder unicode map@ shows the tree that +-- implements the map. If @inorder@ is true, the tree is arranged so that keys +-- will appear in their natural order; otherwise, parents will always precede +-- their children. If @unicode@ is true, the tree will be drawn with Unicode +-- Box Drawing characters; otherwise, the tree will be drawn in ASCII art. +showTreeWith :: Show a => Bool -> Bool -> IntMap a -> String +showTreeWith inorder unicode = unlines . start where + start (IntMap Empty) = [] + start (IntMap (NonEmpty min minV node)) = (show (boundKey min) ++ ":=" ++ show minV) : goL True True node + + goL root fromBefore Tip = [(if fromBefore then if root then ur else rtee else dr) ++ tip] + goL root fromBefore (Bin max maxV l r) + | inorder = fmap (indent fromBefore) children ++ [line] + | otherwise = [line] ++ fmap (indent (not root)) children + where + children = goL False (not inorder) l ++ goR (not inorder) r + line = (if fromBefore then if root then ur else rtee else dr) + ++ horiz ++ horiz ++ horiz + ++ (if inorder then ul else dl) + ++ " " ++ show (boundKey max) ++ ":=" ++ show maxV + + goR fromBefore Tip = [(if fromBefore then ur else rtee) ++ tip] + goR fromBefore (Bin min minV l r) = [line] ++ fmap (indent (not fromBefore)) children + where + children = goL False True l ++ goR True r + line = (if fromBefore then ur else rtee) + ++ horiz ++ horiz ++ horiz + ++ dl + ++ " " ++ show (boundKey min) ++ ":=" ++ show minV + + indent lined line = prefix ++ " " ++ line + where + prefix = if lined then vert else " " + + ul = if unicode then "┘" else "'" + dl = if unicode then "┐" else "." + ur = if unicode then "└" else "`" + dr = if unicode then "┌" else "," + rtee = if unicode then "├" else "+" + vert = if unicode then "│" else "|" + horiz = if unicode then "─" else "-" + tip = if unicode then "╼" else "-*" + +-- | /O(n)/. Test if the internal map structure is valid. +-- +-- >>> valid (fromAscList [(3,"b"), (5,"a")]) +-- True +-- +-- >>> valid (fromAscList [(5,"a"), (3,"b")]) +-- False +-- +-- For information about the invariants that are checked, see 'IntMap_'. +valid :: IntMap a -> Bool +valid = validWith const (&&) + +-- | /O(n)/. Test if the internal map structure is valid, returning extra +-- information about why the map is invalid if it is so. To return this info, +-- @'validWith' assert (.&&.)@ passes a description of everything it asserts +-- to @assert@ along with the assertion's trutheness, then combines the +-- results of these assertions with @.&&.@. +validWith :: (Bool -> String -> prop) -> (prop -> prop -> prop) -> IntMap a -> prop +validWith assert (.&&.) = start + where + start (IntMap Empty) = assert True "Empty maps are always valid." + start (IntMap (NonEmpty _ _ Tip)) = assert True "Singleton maps are always valid." + start (IntMap (NonEmpty min _ (Bin max _ l r))) = + assertInMinBound (boundKey max) min + .&&. goL min max l + .&&. goR min max r + + -- When recursing, we already know that @innerMax < max@, so checking in + -- the subtree that keys are less than @innerMax@ also shows that they + -- are less than @max@. Similarly, we can replace @min@ with @innerMin@ in + -- 'goR'. + goL !_ !_ Tip = assert True "Leaf nodes are always valid." + goL !min !max (Bin innerMax _ l r) = + assertInMinBound (boundKey innerMax) min + .&&. assertInMaxBound (boundKey innerMax) max + .&&. assert (xor (boundKey innerMax) min < xor (boundKey innerMax) max) + (trieError min max (boundKey innerMax) True) + .&&. goL min innerMax l + .&&. goR min innerMax r + + goR !_ !_ Tip = assert True "Leaf nodes are always valid." + goR !min !max (Bin innerMin _ l r) = + assertInMinBound (boundKey innerMin) min + .&&. assertInMaxBound (boundKey innerMin) max + .&&. assert (xor (boundKey innerMin) min > xor (boundKey innerMin) max) + (trieError min max (boundKey innerMin) False) + .&&. goL innerMin max l + .&&. goR innerMin max r + + assertInMinBound k min = assert (compareMinBound k min == InBound) ("Ordering invariant: expected key " ++ show k ++ " > minimum bound " ++ show (boundKey min)) + assertInMaxBound k max = assert (compareMaxBound k max == InBound) ("Ordering invariant: expected key " ++ show k ++ " < maximum bound " ++ show (boundKey max)) + + showBinary k = showIntAtBase 2 intToDigit (fromIntegral k :: Word) "" + + trieError min max k isLeft = "Trie invariant: between " ++ show (boundKey min) ++ " and " ++ show (boundKey max) + ++ ", " ++ show k ++ " was expected to share more bits with " ++ show (if isLeft then boundKey min else boundKey max) + ++ " as it is on the " ++ (if isLeft then "left" else "right") ++ " branch:" + ++ "\n min: " ++ replicate (binLength - length binMin) '0' ++ binMin + ++ "\n k: " ++ replicate (binLength - length binK) '0' ++ binK + ++ "\n max: " ++ replicate (binLength - length binMax) '0' ++ binMax + where + binMin = showBinary (boundKey min) + binK = showBinary k + binMax = showBinary (boundKey max) + binLength = maximum [length binMin, length binK, length binMax] diff --git a/containers/src/Data/IntMap/Lazy.hs b/containers/src/Data/IntMap/Lazy.hs index 8d4ce8c82..1fe25e6c3 100644 --- a/containers/src/Data/IntMap/Lazy.hs +++ b/containers/src/Data/IntMap/Lazy.hs @@ -1,15 +1,21 @@ -{-# LANGUAGE CPP #-} +{-# LANGUAGE CPP, BangPatterns #-} + +#include "containers.h" + #if !defined(TESTING) && defined(__GLASGOW_HASKELL__) +#if USE_REWRITE_RULES +{-# LANGUAGE Trustworthy #-} +#else {-# LANGUAGE Safe #-} #endif - -#include "containers.h" +#endif ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Lazy --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 +-- Copyright : Documentation & Interface (c) Daan Leijen 2002 +-- Documentation (c) Andriy Palamarchuk 2008 +-- Documentation & Implementation (c) Jonathan "gereeter" S. 2020 -- License : BSD-style -- Maintainer : libraries@haskell.org -- Portability : portable @@ -50,29 +56,17 @@ -- -- == 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. --- +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). ----------------------------------------------------------------------------- module Data.IntMap.Lazy ( -- * Map type -#if !defined(TESTING) - IntMap, Key -- instance Eq,Show -#else - IntMap(..), Key -- instance Eq,Show -#endif + IntMap, Key -- * Construction , empty @@ -108,7 +102,7 @@ module Data.IntMap.Lazy ( -- * Query -- ** Lookup - , IM.lookup + , lookup , (!?) , (!) , findWithDefault @@ -120,11 +114,10 @@ module Data.IntMap.Lazy ( , lookupGE -- ** Size - , IM.null + , null , size -- * Combine - -- ** Union , union , unionWith @@ -146,12 +139,12 @@ module Data.IntMap.Lazy ( -- ** Disjoint , disjoint - -- ** Universal combining function + -- ** Deprecated, unsafe general combining function , mergeWithKey -- * Traversal -- ** Map - , IM.map + , map , mapWithKey , traverseWithKey , mapAccum @@ -162,8 +155,8 @@ module Data.IntMap.Lazy ( , mapKeysMonotonic -- * Folds - , IM.foldr - , IM.foldl + , foldr + , foldl , foldrWithKey , foldlWithKey , foldMapWithKey @@ -183,30 +176,30 @@ module Data.IntMap.Lazy ( -- ** Lists , toList - -- ** Ordered lists + -- ** Ordered Lists , toAscList , toDescList -- * Filter - , IM.filter + , filter , filterWithKey , restrictKeys , withoutKeys , partition , partitionWithKey - , mapMaybe , mapMaybeWithKey , mapEither , mapEitherWithKey - , split , splitLookup , splitRoot -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy + , isSubmapOf + , isSubmapOfBy + , isProperSubmapOf + , isProperSubmapOfBy -- * Min\/Max , lookupMin @@ -226,14 +219,1191 @@ module Data.IntMap.Lazy ( , minViewWithKey , maxViewWithKey -#ifdef __GLASGOW_HASKELL__ +#if defined(__GLASGOW_HASKELL__) -- * Debugging , showTree , showTreeWith #endif - ) where +) where -import Data.IntMap.Internal as IM hiding (showTree, showTreeWith) -#ifdef __GLASGOW_HASKELL__ +import Data.IntMap.Internal +import qualified Data.IntMap.Merge.Internal as Merge (runWhenMissingAll) +import qualified Data.IntMap.Merge.Lazy as Merge (merge, mapMissing, mapMaybeMissing, traverseMissing, zipWithMaybeMatched) +#if defined(__GLASGOW_HASKELL__) import Data.IntMap.Internal.DeprecatedDebug #endif + +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (runIdentity) +#else +import Data.IntMap.Merge.Internal (runIdentity) +import Control.Applicative (Applicative(..), (<$>)) +#endif + +import Utils.Containers.Internal.StrictPair (StrictPair(..), toPair) + +import qualified Data.Foldable (Foldable, foldl') +import qualified Data.List (foldl', map) +import qualified Data.IntSet (IntSet, toList) + +import Prelude hiding (foldr, foldl, lookup, null, map, filter, min, max) + +noeval :: a -> () +noeval _ = () + +-- | /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 v = IntMap (NonEmpty (Bound k) v Tip) + +-- | /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 = insertLazy + +-- | /O(min(n,W))/. Insert with a combining function. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert @f new_value old_value@. +-- +-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] +-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" +insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWith = insertWithEval noeval + +-- | /O(min(n,W))/. Insert with a combining function. +-- @'insertWithKey' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert @f key new_value old_value@. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] +-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" +insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithKey f k = insertWith (f k) k + +-- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- and the second element equal to (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) +-- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) +-- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) +insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) +insertLookupWithKey combine !k !v = toPair . start + where + start (IntMap Empty) = Nothing :*: IntMap (NonEmpty (Bound k) v Tip) + start (IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> let mv :*: root' = goL (xor k min) min root + in mv :*: IntMap (NonEmpty min minV root') + OutOfBound -> Nothing :*: IntMap (NonEmpty (Bound k) v (insertMinL (xor k min) min minV root)) + Matched -> Just minV :*: IntMap (NonEmpty (Bound k) (combine k v minV) root) + + goL !_ _ Tip = Nothing :*: Bin (Bound k) v Tip Tip + goL !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> let mv :*: l' = goL xorCache min l + in mv :*: Bin max maxV l' r + | otherwise -> let mv :*: r' = goR xorCacheMax max r + in mv :*: Bin max maxV l r' + OutOfBound | xor (boundKey max) min < xorCacheMax -> Nothing :*: Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Nothing :*: Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Just maxV :*: Bin max (combine k v maxV) l r + where xorCacheMax = xor k max + + goR !_ _ Tip = Nothing :*: Bin (Bound k) v Tip Tip + goR !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> let mv :*: r' = goR xorCache max r + in mv :*: Bin min minV l r' + | otherwise -> let mv :*: l' = goL xorCacheMin min l + in mv :*: Bin min minV l' r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Nothing :*: Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Nothing :*: Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Just minV :*: Bin min (combine k v minV) l r + where xorCacheMin = xor k min + +-- | /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 = start + where + start (IntMap Empty) = IntMap Empty + start m@(IntMap (NonEmpty min minV node)) = case compareMinBound k min of + InBound -> IntMap (NonEmpty min minV (goL (xor k min) node)) + OutOfBound -> m + Matched -> IntMap (NonEmpty min (f minV) node) + + goL !_ Tip = Tip + goL !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goL xorCache l) r + | otherwise -> Bin max maxV l (goR xorCacheMax r) + OutOfBound -> n + Matched -> Bin max (f maxV) l r + where xorCacheMax = xor k max + + goR !_ Tip = Tip + goR !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goR xorCache r) + | otherwise -> Bin min minV (goL xorCacheMin l) r + OutOfBound -> n + Matched -> Bin min (f minV) l r + where xorCacheMin = xor k min + +-- | /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 = adjust (f k) k + +-- | /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 !k = start + where + start (IntMap Empty) = IntMap Empty + start m@(IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> IntMap (NonEmpty min minV (goL (xor k min) root)) + OutOfBound -> m + Matched -> case f minV of + Nothing -> IntMap (nodeToMapL root) + Just minV' -> IntMap (NonEmpty min minV' root) + + goL !_ Tip = Tip + goL !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goL xorCache l) r + | otherwise -> Bin max maxV l (goR xorCacheMax r) + OutOfBound -> n + Matched -> case f maxV of + Nothing -> extractBinL l r + Just maxV' -> Bin max maxV' l r + where xorCacheMax = xor k max + + goR !_ Tip = Tip + goR !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goR xorCache r) + | otherwise -> Bin min minV (goL xorCacheMin l) r + OutOfBound -> n + Matched -> case f minV of + Nothing -> extractBinR l r + Just minV' -> Bin min minV' l r + where xorCacheMin = xor k min + +-- | /O(min(n,W))/. The expression (@'updateWithKey' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a +updateWithKey f k = update (f k) k + +-- | /O(min(n,W))/. Lookup and update. +-- The function returns original value, if it is updated. +-- This is different behavior than 'Data.Map.updateLookupWithKey'. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") +updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a, IntMap a) +updateLookupWithKey = start + where + start _ !_ (IntMap Empty) = (Nothing, IntMap Empty) + start f !k m@(IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> let mv :*: root' = goL f k (xor k min) root + in (mv, IntMap (NonEmpty min minV root')) + OutOfBound -> (Nothing, m) + Matched -> (Just minV, case f (boundKey min) minV of + Nothing -> IntMap (nodeToMapL root) + Just minV' -> IntMap (NonEmpty min minV' root)) + + goL _ !_ !_ Tip = Nothing :*: Tip + goL f !k !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> let mv :*: l' = goL f k xorCache l + in mv :*: Bin max maxV l' r + | otherwise -> let mv :*: r' = goR f k xorCacheMax r + in mv :*: Bin max maxV l r' + OutOfBound -> Nothing :*: n + Matched -> Just maxV :*: case f (boundKey max) maxV of + Nothing -> extractBinL l r + Just maxV' -> Bin max maxV' l r + where xorCacheMax = xor k max + + goR _ !_ !_ Tip = Nothing :*: Tip + goR f !k !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> let mv :*: r' = goR f k xorCache r + in mv :*: Bin min minV l r' + | otherwise -> let mv :*: l' = goL f k xorCacheMin l + in mv :*: Bin min minV l' r + OutOfBound -> Nothing :*: n + Matched -> Just minV :*: case f (boundKey min) minV of + Nothing -> extractBinR l r + Just minV' -> Bin min minV' l r + where xorCacheMin = xor k min + +-- | /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 m = case lookup k m of + Nothing -> case f Nothing of + Nothing -> m + Just v -> insert k v m + Just v -> case f (Just v) of + Nothing -> delete k m + Just v' -> insert k v' m + +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, +-- or update a value in an 'IntMap'. In short : @'lookup' k '<$>' 'alterF' f k m = f +-- ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing -> do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) -> do +-- putStrLn "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserresponse2 :: IO (Maybe String) +-- @ +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- 'Control.Lens.At'. +-- +-- @since 0.5.8 +alterF :: Functor f => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) +alterF f k m = case lookup k m of + Nothing -> fmap (\ret -> case ret of + Nothing -> m + Just v -> insert k v m) (f Nothing) + Just v -> fmap (\ret -> case ret of + Nothing -> delete k m + Just v' -> insert k v' m) (f (Just v)) + +-- | /O(n+m)/. The union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] +unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWith f = unionWithKey (const f) + +-- | /O(n+m)/. The union with a combining function. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] +{-# INLINE unionWithKey #-} +unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithKey combine = unionWithUKey (\k a b -> combine (box k) a b) + +-- | /O(n+m)/. The union with a combining function taking an unboxed key. +-- Identical in functionality to 'unionWithKey'. +unionWithUKey :: (UKey -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithUKey = start + where + start _ (IntMap Empty) m2 = m2 + start _ m1 (IntMap Empty) = m1 + start combine (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (NonEmpty min1 minV1 (goL2 combine minV2 min1 root1 min2 root2)) + | min1 > min2 = IntMap (NonEmpty min2 minV2 (goL1 combine minV1 min1 root1 min2 root2)) + | otherwise = IntMap (NonEmpty min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 root1 root2)) -- we choose min1 arbitrarily, as min1 == min2 + + goL1 _ minV1 !min1 Tip !_ Tip = Bin (minToMax min1) minV1 Tip Tip + goL1 _ minV1 !min1 !n1 !min2 Tip = insertMinL (xor (boundKey min1) min2) min1 minV1 n1 + goL1 _ minV1 !min1 !n1 !min2 n2@(Bin max2 _ _ _) | boundsDisjoint min1 max2 = unionDisjointL minV1 min2 n2 min1 n1 + goL1 combine minV1 !min1 Tip !min2 !n2 = goInsertL1 combine (boundKey min1) minV1 (xor (boundKey min1) min2) min2 n2 + goL1 combine minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> Bin max2 maxV2 (goL1 combine minV1 min1 n1 min2 l2) r2 -- we choose min1 arbitrarily - we just need something from tree 1 + | max1 > max2 -> Bin max1 maxV1 l2 (goR2 combine maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> Bin max2 maxV2 l2 (goR1 combine maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2) + | otherwise -> Bin max1 (combine (boundUKey max1) maxV1 maxV2) l2 (goRFused combine max1 (Bin min1 minV1 l1 r1) r2) -- we choose max1 arbitrarily, as max1 == max2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 (combine (boundUKey max1) maxV1 maxV2) (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> Bin max1 maxV1 (goL1 combine minV1 min1 l1 min2 n2) r1 + + goL2 _ minV2 !_ Tip !min2 Tip = Bin (minToMax min2) minV2 Tip Tip + goL2 _ minV2 !min1 Tip !min2 !n2 = insertMinL (xor (boundKey min2) min1) min2 minV2 n2 + goL2 _ minV2 !min1 n1@(Bin max1 _ _ _) !min2 !n2 | boundsDisjoint min2 max1 = unionDisjointL minV2 min1 n1 min2 n2 + goL2 combine minV2 !min1 !n1 !min2 Tip = goInsertL2 combine (boundKey min2) minV2 (xor (boundKey min2) min1) min1 n1 + goL2 combine minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 n2) r1 -- we choose min2 arbitrarily - we just need something from tree 2 + | max1 > max2 -> Bin max1 maxV1 l1 (goR2 combine maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | max1 < max2 -> Bin max2 maxV2 l1 (goR1 combine maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | otherwise -> Bin max1 (combine (boundUKey max1) maxV1 maxV2) l1 (goRFused combine max1 r1 (Bin min2 minV2 l2 r2)) -- we choose max1 arbitrarily, as max1 == max2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 (combine (boundUKey max1) maxV1 maxV2) (goL2 combine minV2 min1 l1 min2 l2) (goRFused combine max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + LT -> Bin max2 maxV2 (goL2 combine minV2 min1 n1 min2 l2) r2 + + -- 'goLFused' is called instead of 'goL' if the minimums of the two trees are the same + -- Note that because of this property, the trees cannot be disjoint, so we can skip most of the checks in 'goL' + goLFused _ !_ Tip n2 = n2 + goLFused _ !_ n1 Tip = n1 + goLFused combine min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> Bin max2 maxV2 (goLFused combine min n1 l2) r2 + EQ | max1 > max2 -> Bin max1 maxV1 (goLFused combine min l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goLFused combine min l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 (combine (boundUKey max1) maxV1 maxV2) (goLFused combine min l1 l2) (goRFused combine max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> Bin max1 maxV1 (goLFused combine min l1 n2) r1 + + goR1 _ maxV1 !max1 Tip !_ Tip = Bin (maxToMin max1) maxV1 Tip Tip + goR1 _ maxV1 !max1 !n1 !max2 Tip = insertMaxR (xor (boundKey max1) max2) max1 maxV1 n1 + goR1 _ maxV1 !max1 !n1 !max2 n2@(Bin min2 _ _ _) | boundsDisjoint min2 max1 = unionDisjointR maxV1 max1 n1 max2 n2 + goR1 combine maxV1 !max1 Tip !max2 !n2 = goInsertR1 combine (boundKey max1) maxV1 (xor (boundKey max1) max2) max2 n2 + goR1 combine maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> Bin min2 minV2 l2 (goR1 combine maxV1 max1 n1 max2 r2) -- we choose max1 arbitrarily - we just need something from tree 1 + | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | otherwise -> Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 (Bin max1 maxV1 l1 r1) l2) r2 -- we choose min1 arbitrarily, as min1 == min2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) -- we choose min1 arbitrarily, as min1 == min2 + GT -> Bin min1 minV1 l1 (goR1 combine maxV1 max1 r1 max2 n2) + + goR2 _ maxV2 !_ Tip !max2 Tip = Bin (maxToMin max2) maxV2 Tip Tip + goR2 _ maxV2 !max1 Tip !max2 !n2 = insertMaxR (xor (boundKey max2) max1) max2 maxV2 n2 + goR2 _ maxV2 !max1 n1@(Bin min1 _ _ _) !max2 !n2 | boundsDisjoint min1 max2 = unionDisjointR maxV2 max2 n2 max1 n1 + goR2 combine maxV2 !max1 !n1 !max2 Tip = goInsertR2 combine (boundKey max2) maxV2 (xor (boundKey max2) max1) max1 n1 + goR2 combine maxV2 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> Bin min1 minV1 l1 (goR2 combine maxV2 max1 r1 max2 n2) -- we choose max2 arbitrarily - we just need something from tree 2 + | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | otherwise -> Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 (Bin max2 maxV2 l2 r2)) r1 -- we choose min1 arbitrarily, as min1 == min2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | otherwise -> Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) -- we choose min1 arbitrarily, as min1 == min2 + LT -> Bin min2 minV2 l2 (goR2 combine maxV2 max1 n1 max2 r2) + + -- 'goRFused' is called instead of 'goR' if the minimums of the two trees are the same + -- Note that because of this property, the trees cannot be disjoint, so we can skip most of the checks in 'goR' + goRFused _ !_ Tip n2 = n2 + goRFused _ !_ n1 Tip = n1 + goRFused combine max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> Bin min2 minV2 l2 (goRFused combine max n1 r2) + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goRFused combine max r1 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max r1 r2) + | otherwise -> Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2) (goRFused combine max r1 r2) -- we choose min1 arbitrarily, as min1 == min2 + GT -> Bin min1 minV1 l1 (goRFused combine max r1 n2) + + goInsertL1 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertL1 combine k v !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goInsertL1 combine k v xorCache min l) r + | otherwise -> Bin max maxV l (goInsertR1 combine k v xorCacheMax max r) + OutOfBound | xor (boundKey max) min < xorCacheMax -> Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Bin max (combine (unbox k) v maxV) l r + where xorCacheMax = xor k max + + goInsertR1 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertR1 combine k v !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goInsertR1 combine k v xorCache max r) + | otherwise -> Bin min minV (goInsertL1 combine k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Bin min (combine (unbox k) v minV) l r + where xorCacheMin = xor k min + + goInsertL2 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertL2 combine k v !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goInsertL2 combine k v xorCache min l) r + | otherwise -> Bin max maxV l (goInsertR2 combine k v xorCacheMax max r) + OutOfBound | xor (boundKey max) min < xorCacheMax -> Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Bin max (combine (unbox k) maxV v) l r + where xorCacheMax = xor k max + + goInsertR2 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertR2 combine k v !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goInsertR2 combine k v xorCache max r) + | otherwise -> Bin min minV (goInsertL2 combine k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Bin min (combine (unbox k) minV v) l r + where xorCacheMin = xor k min + +-- | 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 :: Data.Foldable.Foldable f => (a -> a -> a) -> f (IntMap a) -> IntMap a +unionsWith f = Data.Foldable.foldl' (unionWith f) empty + +-- | /O(n+m)/. Difference with a combining function. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) +-- > == singleton 3 "b:B" +differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWith f = differenceWithKey (const f) + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). +-- If it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) +-- > == singleton 3 "3:b|B" +{-# INLINE differenceWithKey #-} +differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWithKey combine = differenceWithUKey (\k a b -> combine (box k) a b) + +-- | /O(n+m)/. The difference with a combining function taking an unboxed +-- key. Identical in functionality to 'differenceWithKey'. +differenceWithUKey :: (UKey -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWithUKey = start + where + start _ (IntMap Empty) !_ = IntMap Empty + start _ !m (IntMap Empty) = m + start combine (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (NonEmpty min1 minV1 (goL2 combine minV2 min1 root1 min2 root2)) + | min1 > min2 = IntMap (goL1 combine minV1 min1 root1 min2 root2) + | otherwise = case combine (boundUKey min1) minV1 minV2 of + Nothing -> IntMap (goLFused combine min1 root1 root2) + Just minV1' -> IntMap (NonEmpty min1 minV1' (goLFusedKeep combine min1 root1 root2)) + + goL1 combine minV1 !min1 Tip !min2 !n2 = goLookupL combine (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL1 _ minV1 !min1 !n1 !_ Tip = NonEmpty min1 minV1 n1 + goL1 _ minV1 !min1 n1@(Bin _ _ _ _) !_ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = NonEmpty min1 minV1 n1 + goL1 combine minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL1 combine minV1 min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> r2lMap $ NonEmpty max1 maxV1 (goR2 combine maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> r2lMap $ goR1 combine maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> r2lMap $ goRFused combine max1 (Bin min1 minV1 l1 r1) r2 + Just maxV1' -> r2lMap $ NonEmpty max1 maxV1' (goRFusedKeep combine max1 (Bin min1 minV1 l1 r1) r2) + EQ | max1 > max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (NonEmpty max1 maxV1 (goR2 combine maxV2 max1 r1 max2 r2)) + | max1 < max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> binL (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max1 r1 r2) + Just maxV1' -> binL (goL1 combine minV1 min1 l1 min2 l2) (NonEmpty max1 maxV1' (goRFusedKeep combine max1 r1 r2)) + GT -> binL (goL1 combine minV1 min1 l1 min2 n2) (NonEmpty max1 maxV1 r1) + + goL2 _ _ !_ Tip !_ !_ = Tip + goL2 combine minV2 !min1 !n1 !min2 Tip = goDeleteL combine (boundKey min2) minV2 (xor (boundKey min2) min1) n1 + goL2 _ _ !_ n1@(Bin max1 _ _ _) !min2 (Bin _ _ _ _) | boundsDisjoint min2 max1 = n1 + goL2 combine minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goL2 combine minV2 min1 n1 min2 l2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> goL2 combine minV2 min1 l1 min2 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goL2 combine minV2 min1 l1 min2 l2) r' + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> case goRFused combine max1 r1 r2 of + Empty -> goL2 combine minV2 min1 l1 min2 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goL2 combine minV2 min1 l1 min2 l2) r' + Just maxV1' -> Bin max1 maxV1' (goL2 combine minV2 min1 l1 min2 l2) (goRFusedKeep combine max1 r1 r2) + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 n2) r1 -- min2 is arbitrary here - we just need something from tree 2 + | max1 > max2 -> Bin max1 maxV1 l1 (goR2 combine maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2) of + Empty -> l1 + NonEmpty max' maxV' r' -> Bin max' maxV' l1 r' + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> case goRFused combine max1 r1 (Bin min2 minV2 l2 r2) of + Empty -> l1 + NonEmpty max' maxV' r' -> Bin max' maxV' l1 r' + Just maxV1' -> Bin max1 maxV1' l1 (goRFusedKeep combine max1 r1 (Bin min2 minV2 l2 r2)) + + goLFused _ !_ Tip !_ = Empty + goLFused _ !_ (Bin max1 maxV1 l1 r1) Tip = case deleteMinL max1 maxV1 l1 r1 of + NE min' minV' n' -> NonEmpty min' minV' n' + goLFused combine !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFused combine min n1 l2 + EQ | max1 > max2 -> binL (goLFused combine min l1 l2) (NonEmpty max1 maxV1 (goR2 combine maxV2 max1 r1 max2 r2)) + | max1 < max2 -> binL (goLFused combine min l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> binL (goLFused combine min l1 l2) (goRFused combine max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + Just maxV1' -> binL (goLFused combine min l1 l2) (NonEmpty max1 maxV1' (goRFusedKeep combine max1 r1 r2)) + GT -> binL (goLFused combine min l1 n2) (NonEmpty max1 maxV1 r1) + + goLFusedKeep _ !_ !n1 Tip = n1 + goLFusedKeep _ !_ Tip !_ = Tip + goLFusedKeep combine !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFusedKeep combine min n1 l2 + EQ | max1 > max2 -> Bin max1 maxV1 (goLFusedKeep combine min l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> goLFusedKeep combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFusedKeep combine min l1 l2) r' + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> case goRFused combine max1 r1 r2 of -- we choose max1 arbitrarily, as max1 == max2 + Empty -> goLFusedKeep combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFusedKeep combine min l1 l2) r' + Just maxV1' -> Bin max1 maxV1' (goLFusedKeep combine min l1 l2) (goRFusedKeep combine max1 r1 r2) + GT -> Bin max1 maxV1 (goLFusedKeep combine min l1 n2) r1 + + goR1 combine maxV1 !max1 Tip !max2 !n2 = goLookupR combine (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR1 _ maxV1 !max1 !n1 !_ Tip = NonEmpty max1 maxV1 n1 + goR1 _ maxV1 !max1 n1@(Bin _ _ _ _) !_ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = NonEmpty max1 maxV1 n1 + goR1 combine maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR1 combine maxV1 max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> l2rMap $ NonEmpty min1 minV1 (goL2 combine minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) + | min1 > min2 -> l2rMap $ goL1 combine minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> l2rMap $ goLFused combine min1 (Bin max1 maxV1 l1 r1) l2 + Just minV1' -> l2rMap $ NonEmpty min1 minV1' (goLFusedKeep combine min1 (Bin max1 maxV1 l1 r1) l2) + EQ | min1 < min2 -> binR (NonEmpty min1 minV1 (goL2 combine minV2 min1 l1 min2 l2)) (goR1 combine maxV1 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> binR (goLFused combine min1 l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) + Just minV1' -> binR (NonEmpty min1 minV1' (goLFusedKeep combine min1 l1 l2)) (goR1 combine maxV1 max1 r1 max2 r2) + GT -> binR (NonEmpty min1 minV1 l1) (goR1 combine maxV1 max1 r1 max2 n2) + + goR2 _ _ !_ Tip !_ !_ = Tip + goR2 combine maxV2 !max1 !n1 !max2 Tip = goDeleteR combine (boundKey max2) maxV2 (xor (boundKey max2) max1) n1 + goR2 _ _ !_ n1@(Bin min1 _ _ _) !max2 (Bin _ _ _ _) | boundsDisjoint min1 max2 = n1 + goR2 combine maxV2 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goR2 combine maxV2 max1 n1 max2 r2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> goR2 combine maxV2 max1 r1 max2 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goR2 combine maxV2 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> case goLFused combine min1 l1 l2 of + Empty -> goR2 combine maxV2 max1 r1 max2 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goR2 combine maxV2 max1 r1 max2 r2) + Just minV1' -> Bin min1 minV1' (goLFusedKeep combine min1 l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> Bin min1 minV1 l1 (goR2 combine maxV2 max1 r1 max2 n2) -- max2 is arbitrary here - we just need something from tree 2 + | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2) of + Empty -> r1 + NonEmpty min' minV' l' -> Bin min' minV' l' r1 + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> case goLFused combine min1 l1 (Bin max2 maxV2 l2 r2) of + Empty -> r1 + NonEmpty min' minV' l' -> Bin min' minV' l' r1 + Just minV1' -> Bin min1 minV1' (goLFusedKeep combine min1 l1 (Bin max2 maxV2 l2 r2)) r1 + + goRFused _ !_ Tip !_ = Empty + goRFused _ !_ (Bin min1 minV1 l1 r1) Tip = case deleteMaxR min1 minV1 l1 r1 of + NE max' maxV' n' -> NonEmpty max' maxV' n' + goRFused combine !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFused combine max n1 r2 + EQ | min1 < min2 -> binR (NonEmpty min1 minV1 (goL2 combine minV2 min1 l1 min2 l2)) (goRFused combine max r1 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max r1 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> binR (goLFused combine min1 l1 l2) (goRFused combine max r1 r2) -- we choose min1 arbitrarily, as min1 == min2 + Just minV1' -> binR (NonEmpty min1 minV1' (goLFusedKeep combine min1 l1 l2)) (goRFused combine max r1 r2) + GT -> binR (NonEmpty min1 minV1 l1) (goRFused combine max r1 n2) + + goRFusedKeep _ !_ !n1 Tip = n1 + goRFusedKeep _ !_ Tip !_ = Tip + goRFusedKeep combine !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFusedKeep combine max n1 r2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goRFusedKeep combine max r1 r2) + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> goRFusedKeep combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFusedKeep combine max r1 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of -- we choose min1 arbitrarily, as min1 == min2 + Nothing -> case goLFused combine min1 l1 l2 of + Empty -> goRFusedKeep combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFusedKeep combine max r1 r2) + Just minV1' -> Bin min1 minV1' (goLFusedKeep combine min1 l1 l2) (goRFusedKeep combine max r1 r2) + GT -> Bin min1 minV1 l1 (goRFusedKeep combine max r1 n2) + + goLookupL _ !k v !_ Tip = NonEmpty (Bound k) v Tip + goLookupL combine !k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL combine k v xorCache l + | otherwise -> goLookupR combine k v xorCacheMax r + OutOfBound -> NonEmpty (Bound k) v Tip + Matched -> case combine (unbox k) v maxV of + Nothing -> Empty + Just v' -> NonEmpty (Bound k) v' Tip + where xorCacheMax = xor k max + + goLookupR _ !k v !_ Tip = NonEmpty (Bound k) v Tip + goLookupR combine !k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR combine k v xorCache r + | otherwise -> goLookupL combine k v xorCacheMin l + OutOfBound -> NonEmpty (Bound k) v Tip + Matched -> case combine (unbox k) v minV of + Nothing -> Empty + Just v' -> NonEmpty (Bound k) v' Tip + where xorCacheMin = xor k min + + goDeleteL _ !_ _ !_ Tip = Tip + goDeleteL combine !k v !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goDeleteL combine k v xorCache l) r + | otherwise -> Bin max maxV l (goDeleteR combine k v xorCacheMax r) + OutOfBound -> n + Matched -> case combine (unbox k) maxV v of + Nothing -> extractBinL l r + Just maxV' -> Bin max maxV' l r + where xorCacheMax = xor k max + + goDeleteR _ !_ _ !_ Tip = Tip + goDeleteR combine !k v !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goDeleteR combine k v xorCache r) + | otherwise -> Bin min minV (goDeleteL combine k v xorCacheMin l) r + OutOfBound -> n + Matched -> case combine (unbox k) minV v of + Nothing -> extractBinR l r + Just minV' -> Bin min minV' l r + where xorCacheMin = xor k min + +-- | /O(n+m)/. The intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" +intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWith f = intersectionWithKey (const f) + +-- | /O(n+m)/. The intersection with a combining function. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" +{-# INLINE intersectionWithKey #-} +intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWithKey combine = intersectionWithUKey (\k a b -> combine (box k) a b) + +-- | /O(n+m)/. The intersection with a combining function taking an unboxed +-- key. Identical in functionality to 'intersectionWithKey'. +intersectionWithUKey :: (UKey -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWithUKey = start + where + start _ (IntMap Empty) !_ = IntMap Empty + start _ !_ (IntMap Empty) = IntMap Empty + start combine (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (goL2 combine minV2 min1 root1 min2 root2) + | min1 > min2 = IntMap (goL1 combine minV1 min1 root1 min2 root2) + | otherwise = IntMap (NonEmpty min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 root1 root2)) -- we choose min1 arbitrarily, as min1 == min2 + + -- TODO: This scheme might produce lots of unnecessary l2r and r2l calls. This should be rectified. + + goL1 _ _ !_ !_ !_ Tip = Empty + goL1 combine minV1 min1 Tip min2 n2 = goLookupL1 combine (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL1 _ _ min1 (Bin _ _ _ _) _ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = Empty + goL1 combine minV1 min1 n1@(Bin max1 maxV1 l1 r1) min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL1 combine minV1 min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> r2lMap $ goR2 combine maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2 + | max1 < max2 -> r2lMap $ goR1 combine maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + | otherwise -> r2lMap $ NonEmpty max1 (combine (boundUKey max1) maxV1 maxV2) (goRFused combine max1 (Bin min1 minV1 l1 r1) r2) + EQ | max1 > max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> r2lMap (NonEmpty max1 (combine (boundUKey max1) maxV1 maxV2) (goRFused combine max1 r1 r2)) + NonEmpty min' minV' l' -> NonEmpty min' minV' (Bin max1 (combine (boundUKey max1) maxV1 maxV2) l' (goRFused combine max1 r1 r2)) + GT -> goL1 combine minV1 min1 l1 min2 n2 + + goL2 _ _ !_ Tip !_ !_ = Empty + goL2 combine minV2 min1 n1 min2 Tip = goLookupL2 combine (boundKey min2) minV2 (xor (boundKey min2) min1) n1 + goL2 _ _ _ (Bin max1 _ _ _) min2 (Bin _ _ _ _) | boundsDisjoint min2 max1 = Empty + goL2 combine minV2 min1 n1@(Bin max1 maxV1 l1 r1) min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goL2 combine minV2 min1 n1 min2 l2 + EQ | max1 > max2 -> binL (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> binL (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case goL2 combine minV2 min1 l1 min2 l2 of + Empty -> r2lMap (NonEmpty max1 (combine (boundUKey max1) maxV1 maxV2) (goRFused combine max1 r1 r2)) + NonEmpty min' minV' l' -> NonEmpty min' minV' (Bin max1 (combine (boundUKey max1) maxV1 maxV2) l' (goRFused combine max1 r1 r2)) + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> goL2 combine minV2 min1 l1 min2 n2 -- min2 is arbitrary here - we just need something from tree 2 + | max1 > max2 -> r2lMap $ goR2 combine maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2) + | max1 < max2 -> r2lMap $ goR1 combine maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2) + | otherwise -> r2lMap $ NonEmpty max1 (combine (boundUKey max1) maxV1 maxV2) (goRFused combine max1 r1 (Bin min2 minV2 l2 r2)) + + goLFused _ !_ Tip !_ = Tip + goLFused _ !_ !_ Tip = Tip + goLFused combine !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFused combine min n1 l2 + EQ | max1 > max2 -> case goR2 combine maxV2 max1 r1 max2 r2 of + Empty -> goLFused combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFused combine min l1 l2) r' + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> goLFused combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFused combine min l1 l2) r' + | otherwise -> Bin max1 (combine (boundUKey max1) maxV1 maxV2) (goLFused combine min l1 l2) (goRFused combine max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> goLFused combine min l1 n2 + + goR1 _ _ !_ !_ !_ Tip = Empty + goR1 combine maxV1 max1 Tip max2 n2 = goLookupR1 combine (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR1 _ _ max1 (Bin _ _ _ _) _ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = Empty + goR1 combine maxV1 max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR1 combine maxV1 max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> l2rMap $ goL2 combine minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | min1 > min2 -> l2rMap $ goL1 combine minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | otherwise -> l2rMap $ NonEmpty min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 (Bin max1 maxV1 l1 r1) l2) + EQ | min1 < min2 -> binR (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> l2rMap (NonEmpty min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2)) + NonEmpty max' maxV' r' -> NonEmpty max' maxV' (Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2) r') + GT -> goR1 combine maxV1 max1 r1 max2 n2 + + goR2 _ _ !_ Tip !_ !_ = Empty + goR2 combine maxV2 max1 n1 max2 Tip = goLookupR2 combine (boundKey max2) maxV2 (xor (boundKey max2) max1) n1 + goR2 _ _ _ (Bin min1 _ _ _) max2 (Bin _ _ _ _) | boundsDisjoint min1 max2 = Empty + goR2 combine maxV2 max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goR2 combine maxV2 max1 n1 max2 r2 + EQ | min1 < min2 -> binR (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | otherwise -> case goR2 combine maxV2 max1 r1 max2 r2 of + Empty -> l2rMap (NonEmpty min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2)) + NonEmpty max' maxV' r' -> NonEmpty max' maxV' (Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2) r') + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> goR2 combine maxV2 max1 r1 max2 n2 -- max2 is arbitrary here - we just need something from tree 2 + | min1 < min2 -> l2rMap $ goL2 combine minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2) + | min1 > min2 -> l2rMap $ goL1 combine minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2) + | otherwise -> l2rMap $ NonEmpty min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 (Bin max2 maxV2 l2 r2)) + + goRFused _ !_ Tip !_ = Tip + goRFused _ !_ !_ Tip = Tip + goRFused combine !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFused combine max n1 r2 + EQ | min1 < min2 -> case goL2 combine minV2 min1 l1 min2 l2 of + Empty -> goRFused combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFused combine max r1 r2) + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> goRFused combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFused combine max r1 r2) + | otherwise -> Bin min1 (combine (boundUKey min1) minV1 minV2) (goLFused combine min1 l1 l2) (goRFused combine max r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + GT -> goRFused combine max r1 n2 + + goLookupL1 _ !_ _ !_ Tip = Empty + goLookupL1 combine k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL1 combine k v xorCache l + | otherwise -> goLookupR1 combine k v xorCacheMax r + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) (combine (unbox k) v maxV) Tip + where xorCacheMax = xor k max + + goLookupR1 _ !_ _ !_ Tip = Empty + goLookupR1 combine k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR1 combine k v xorCache r + | otherwise -> goLookupL1 combine k v xorCacheMin l + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) (combine (unbox k) v minV) Tip + where xorCacheMin = xor k min + + goLookupL2 _ !_ _ !_ Tip = Empty + goLookupL2 combine k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL2 combine k v xorCache l + | otherwise -> goLookupR2 combine k v xorCacheMax r + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) (combine (unbox k) maxV v) Tip + where xorCacheMax = xor k max + + goLookupR2 _ !_ _ !_ Tip = Empty + goLookupR2 combine k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR2 combine k v xorCache r + | otherwise -> goLookupL2 combine k v xorCacheMin l + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) (combine (unbox k) minV v) Tip + where xorCacheMin = xor k min + +-- | /O(n+m)/. An unsafe general combining function. +-- +-- WARNING: This function can produce corrupt maps and its results may depend +-- on the internal structures of its inputs. Users should prefer +-- 'Data.IntMap.Merge.Lazy.merge' or 'Data.IntMap.Merge.Lazy.mergeA'. This +-- function is also significantly slower than 'Data.IntMap.Merge.Lazy.merge'. +-- +-- When 'mergeWithKey' is given three arguments, it is inlined to the call +-- site. You should therefore use 'mergeWithKey' only to define custom +-- combining functions. For example, you could define 'unionWithKey', +-- 'differenceWithKey' and 'intersectionWithKey' as +-- +-- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 +-- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 +-- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 +-- +-- When calling @'mergeWithKey' combine only1 only2@, a function combining two +-- '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@, +-- @'filterWithKey' f@, or @'mapMaybeWithKey' f@ could be used for any @f@. +mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) -> IntMap a -> IntMap b -> IntMap c +mergeWithKey matched miss1 miss2 = Merge.merge (Merge.mapMaybeMissing (single miss1)) (Merge.mapMaybeMissing (single miss2)) (Merge.zipWithMaybeMatched matched) where + single miss k v = case miss (IntMap (NonEmpty (Bound k) v Tip)) of + IntMap Empty -> Nothing + IntMap (NonEmpty _ v' _) -> Just v' + +-- | /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 = mapLazy + +-- | /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 = runIdentity . Merge.runWhenMissingAll (Merge.mapMissing f) + +#if USE_REWRITE_RULES +{-# NOINLINE[1] mapWithKey #-} +{-# RULES +"map/mapWithKey" forall f g m . mapLazy f (mapWithKey g m) = mapWithKey (\k -> f . g k) m +"mapWithKey/map" forall f g m . mapWithKey f (mapLazy g m) = mapWithKey (\k -> f k . g) m +"mapWithKey/mapWithKey" forall f g m . mapWithKey f (mapWithKey g m) = mapWithKey (\k -> f k . g k) m + #-} +#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 +{-# INLINE traverseWithKey #-} +traverseWithKey :: Applicative f => (Key -> a -> f b) -> IntMap a -> f (IntMap b) +traverseWithKey f = Merge.runWhenMissingAll (Merge.traverseMissing f) + +-- | /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 = start + where + start a (IntMap Empty) = (a, IntMap Empty) + start a (IntMap (NonEmpty min minV root)) = + let (a', minV') = f a (boundKey min) minV + (a'', root') = goL root a' + in (a'', IntMap (NonEmpty min minV' root')) + + goL Tip a = (a, Tip) + goL (Bin max maxV l r) a = + let (a', l') = goL l a + (a'', r') = goR r a' + (a''', maxV') = f a'' (boundKey max) maxV + in (a''', Bin max maxV' l' r') + + goR Tip a = (a, Tip) + goR (Bin min minV l r) a = + let (a', minV') = f a (boundKey min) minV + (a'', l') = goL l a' + (a''', r') = goR r a'' + in (a''', Bin min minV' l' r') + +-- | /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 = start + where + start a (IntMap Empty) = (a, IntMap Empty) + start a (IntMap (NonEmpty min minV root)) = + let (a', root') = goL root a + (a'', minV') = f a' (boundKey min) minV + in (a'', IntMap (NonEmpty min minV' root')) + + goL Tip a = (a, Tip) + goL (Bin max maxV l r) a = + let (a', maxV') = f a (boundKey max) maxV + (a'', r') = goR r a' + (a''', l') = goL l a'' + in (a''', Bin max maxV' l' r') + + goR Tip a = (a, Tip) + goR (Bin min minV l r) a = + let (a', r') = goR r a + (a'', l') = goL l a' + (a''', minV') = f a'' (boundKey min) minV + in (a''', Bin min minV' l' r') + +-- | /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 = foldlWithKey' (\m k a -> insert (f k) a m) empty + +-- | /O(n*min(n,W))/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" +mapKeysWith :: (a -> a -> a) -> (Key -> Key) -> IntMap a -> IntMap a +mapKeysWith combine f = foldlWithKey' (\m k a -> insertWith combine (f k) a m) empty + +-- | /O(n*min(n,W))/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has slightly better performance than 'mapKeys'. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] +mapKeysMonotonic :: (Key -> Key) -> IntMap a -> IntMap a +mapKeysMonotonic = mapKeys + +-- | /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) -> Data.IntSet.IntSet -> IntMap a +fromSet f = fromDistinctAscList . Data.List.map (\k -> (k, f k)) . Data.IntSet.toList + +-- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs. +fromList :: [(Key, a)] -> IntMap a +fromList = fromListLazy + +-- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")] +-- > fromListWith (++) [] == empty +fromListWith :: (a -> a -> a) -> [(Key, a)] -> IntMap a +fromListWith f = Data.List.foldl' (\t (k, a) -> insertWith f k a t) empty + +-- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")] +-- > fromListWithKey f [] == empty +fromListWithKey :: (Key -> a -> a -> a) -> [(Key, a)] -> IntMap a +fromListWithKey f = Data.List.foldl' (\t (k, a) -> insertWithKey f k a t) empty + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +fromAscList :: [(Key, a)] -> IntMap a +fromAscList = start where + start [] = IntMap Empty + start ((min, minV) : rest) = IntMap (go min minV rest StackBase) + + go !k v [] !stk = completeBuildStack (Bound k) v Tip stk + go !k v ((next, nextV) : rest) !stk + | next == k = go k nextV rest stk + | otherwise = go next nextV rest (pushBuildStack (xor k (Bound next)) k v Tip stk) + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +fromAscListWith :: (a -> a -> a) -> [(Key, a)] -> IntMap a +fromAscListWith f = fromAscListWithKey (const f) + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order, with a combining function on equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")] +fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key, a)] -> IntMap a +fromAscListWithKey f = start where + start [] = IntMap Empty + start ((min, minV) : rest) = IntMap (go min minV rest StackBase) + + go !k v [] !stk = completeBuildStack (Bound k) v Tip stk + go !k v ((next, nextV) : rest) !stk + | next == k = go k (f k nextV v) rest stk + | otherwise = go next nextV rest (pushBuildStack (xor k (Bound next)) k v Tip stk) + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order and all distinct. +-- /The precondition (input list is strictly ascending) is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +fromDistinctAscList :: [(Key, a)] -> IntMap a +fromDistinctAscList = start where + start [] = IntMap Empty + start ((min, minV) : rest) = IntMap (go min minV rest StackBase) + + go !k v [] !stk = completeBuildStack (Bound k) v Tip stk + go !k v ((next, nextV) : rest) !stk = go next nextV rest (pushBuildStack (xor k (Bound next)) k v Tip stk) + +-- | /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 (const f) + +-- | /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" +{-# INLINE mapMaybeWithKey #-} +mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b +mapMaybeWithKey f = runIdentity . Merge.runWhenMissingAll (Merge.mapMaybeMissing f) + +-- | /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 = mapEitherWithKey (const f) + +-- | /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")]) +{-# INLINE mapEitherWithKey #-} +mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) +mapEitherWithKey f = mapEitherWithUKey (\k a -> f (box k) a) + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results with +-- a mapping function taking unboxed keys. Identical in functionality to +-- 'mapEitherWithKey'. +mapEitherWithUKey :: (UKey -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) +mapEitherWithUKey func = start + where + start (IntMap Empty) = (IntMap Empty, IntMap Empty) + start (IntMap (NonEmpty min minV root)) = case func (boundUKey min) minV of + Left v -> let t :*: f = goTrueL root + in (IntMap (NonEmpty min v t), IntMap f) + Right v -> let t :*: f = goFalseL root + in (IntMap t, IntMap (NonEmpty min v f)) + + goTrueL Tip = Tip :*: Empty + goTrueL (Bin max maxV l r) = case func (boundUKey max) maxV of + Left v -> let tl :*: fl = goTrueL l + tr :*: fr = goTrueR r + in Bin max v tl tr :*: binL fl fr + Right v -> let tl :*: fl = goTrueL l + tr :*: fr = goFalseR r + in binNodeMapL tl tr :*: binL fl (NonEmpty max v fr) + + goTrueR Tip = Tip :*: Empty + goTrueR (Bin min minV l r) = case func (boundUKey min) minV of + Left v -> let tl :*: fl = goTrueL l + tr :*: fr = goTrueR r + in Bin min v tl tr :*: binR fl fr + Right v -> let tl :*: fl = goFalseL l + tr :*: fr = goTrueR r + in binMapNodeR tl tr :*: binR (NonEmpty min v fl) fr + + goFalseL Tip = Empty :*: Tip + goFalseL (Bin max maxV l r) = case func (boundUKey max) maxV of + Left v -> let tl :*: fl = goFalseL l + tr :*: fr = goTrueR r + in binL tl (NonEmpty max v tr) :*: binNodeMapL fl fr + Right v -> let tl :*: fl = goFalseL l + tr :*: fr = goFalseR r + in binL tl tr :*: Bin max v fl fr + + goFalseR Tip = Empty :*: Tip + goFalseR (Bin min minV l r) = case func (boundUKey min) minV of + Left v -> let tl :*: fl = goTrueL l + tr :*: fr = goFalseR r + in binR (NonEmpty min v tl) tr :*: binMapNodeR fl fr + Right v -> let tl :*: fl = goFalseL l + tr :*: fr = goFalseR r + in binR tl tr :*: Bin min v fl fr + +-- | /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))/. 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. +-- +-- > 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 _ (IntMap Empty) = IntMap Empty +updateMinWithKey f (IntMap (NonEmpty min minV root)) = IntMap $ case f (boundKey min) minV of + Nothing -> nodeToMapL root + Just minV' -> NonEmpty min minV' root + +-- | /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 _ (IntMap Empty) = IntMap Empty +updateMaxWithKey f (IntMap (NonEmpty min minV Tip)) = IntMap $ case f (boundKey min) minV of + Nothing -> Empty + Just minV' -> NonEmpty min minV' Tip +updateMaxWithKey f (IntMap (NonEmpty min minV (Bin max maxV l r))) = IntMap . NonEmpty min minV $ case f (boundKey max) maxV of + Nothing -> extractBinL l r + Just maxV' -> Bin max maxV' l r diff --git a/containers/src/Data/IntMap/Merge/Internal.hs b/containers/src/Data/IntMap/Merge/Internal.hs new file mode 100644 index 000000000..f5c85f95f --- /dev/null +++ b/containers/src/Data/IntMap/Merge/Internal.hs @@ -0,0 +1,896 @@ +{-# LANGUAGE CPP, BangPatterns #-} + +#include "containers.h" + +#if !defined(TESTING) && defined(__GLASGOW_HASKELL__) +#if MIN_VERSION_base(4,8,0) || __GLASGOW_HASKELL__ < 708 +{-# LANGUAGE Safe #-} +#else +{-# LANGUAGE Trustworthy #-} +#endif +#endif + +{-# OPTIONS_HADDOCK not-home #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.IntMap.Merge.Internal +-- Copyright : (c) Jonathan "gereeter" S. 2020 +-- License : BSD-style +-- Maintainer : libraries@haskell.org +-- Stability : provisional +-- Portability : portable +-- +-- = WARNING +-- +-- This module is considered __internal__. +-- +-- The Package Versioning Policy __does not apply__. +-- +-- This contents of this module may change __in any way whatsoever__ +-- and __without any warning__ between minor versions of this package. +-- +-- Authors importing this module are expected to track development +-- closely. +-- +-- = Description +-- +-- This defines the data structures and core (hidden) manipulations +-- on representations. +-- +-- = The Merging Process +-- +-- Since keys have consistent positions in trie-like structures, merging is +-- largely a problem of alignment. Considering the sets of keys matching the +-- bit prefix of a node, we have the following six cases: +-- +-- 1. Disjoint ranges +-- +-- > Node 1: ## +-- > Node 2: #### +-- +-- Depending on the type of merge, we combine the nodes. If +-- 'Data.IntMap.Lazy.union', we create a new 'Bin' combining the nodes. +-- If 'Data.IntMap.Lazy.intersection', we return 'Empty'. In general, both +-- nodes get passed through the 'WhenMissing' tactics then combined. +-- +-- 2. Equal ranges +-- +-- > Node 1: #### +-- > Node 2: #### +-- +-- We recurively merge node 1's left branch with node 2's left branch and +-- node 1's right branch with node 2's right branch. +-- +-- 3. 1-in-2 (left) +-- +-- > Node 1: ## +-- > Node 2: #### +-- +-- We recursively merge node 1 with node 2's left branch. +-- +-- 4. 1-in-2 (right) +-- +-- > Node 1: ## +-- > Node 2: #### +-- +-- We recursively merge node 1 with node 2's right branch. +-- +-- 5. 2-in-1 (left) +-- +-- > Node 1: #### +-- > Node 2: ## +-- +-- We recursively merge node 1's left branch with node 2. +-- +-- 6. 2-in-1 (right) +-- +-- > Node 1: #### +-- > Node 2: ## +-- +-- We recursively merge node 1's right branch with node 2. +-- +-- Distinguishing the latter 5 cases is much harder with the disjoint ranges +-- case in the mix, so eliminating that is the first step. With the min/max +-- implicit representation, we can test for that case by seeing if the minimum +-- of one node is greater than the maximum of the other node 'boundsDisjoint'. +-- Technically, this condition can also trigger in the other cases, since the +-- cases defined above are about shared bits, not being in between the minima +-- and maxima. For example, a minimum of -1 shares no bits with a maximum of 0, +-- so for the purposes of the cases above, a node with those bounds would have +-- a "shared bits range" of all integers. However, determining that two nodes +-- will never overlap is a useful condition in its own right. If taking an +-- intersection, for example, we can immediately return 'Empty', and even the +-- union case can be simplified (see 'unionDisjointL' and 'unionDisjointR') if +-- not to a single 'Bin' node. +-- +-- Once the ranges of the nodes are known to be overlapping, we can compare +-- range sizes to distinguish between the equal, 1-in-2, and 2-in-1 cases. When +-- the minumum and maximum bounds for a node are XORed with each other (using +-- 'xorBounds'), all shared prefix bits will cancel each other out and produce +-- zeros, and the first bit where the bounds disagree will become the result's +-- most significant set bit (MSB). This justifies using 'compareMSB' to compare +-- how many shared prefix bits the two nodes have. The left and right variants +-- of the 1-in-2 and 2-in-1 cases can then be distinguished using the same +-- techniques as in single-key queries, taking a key from the smaller node and +-- determining which branch it belongs in. +-- +-- == Bound Complications +-- +-- Unfortunately, since our tree structure is a bit more complicated than a +-- PATRICIA tree, there is more complexity involved. Instead of just aligning +-- internal nodes and merging values at the keys, we need to interleave the +-- combination of values with the alignment of nodes. At every point in the +-- recursion, whenever we produce a composite node, we need to produce the key +-- and value that go along with that node. +-- +-- Determining which bound to keep and what values to combine is +-- straightforward; the lesser of the two minima is the new minimum, and the +-- greater of the two maxima is new maximum. However, the unused key/value +-- pair (if there is one) needs to be pushed downward in the recursion to be +-- included in the merged map. This necessitates three variants of each helper +-- function, corresponding to the three choices of pushing down the value from +-- node 1, pushing down the value from node 2, and pushing down neither. For +-- example, @goL1@ is called when @min1 > min2@, @goL2@ when @min2 > min1@, and +-- @goLFused@ when @min1 = min2@. +-- +-- As we are forced to do these comparisons anyway, we can use them to reduce +-- the number of cases to consider. When @min1 > min2@, for example, @min2@ +-- cannot be greater than @max1@, so a single comparison suffices to determine +-- whether the ranges are disjoint, and the 2-in-1 (right) case is impossible. +-- +-- == Base Case: Merging leaves into trees +-- +-- In the base case, when the recursion hits the leaves, there are two cases. +-- If the 'Tip' corresponds to the bound that was pushed down from higher up +-- in the tree, then the merge operation begins to look like an insertion (in +-- the case of 'Data.IntMap.Lazy.union'), deletion (in +-- 'Data.IntMap.Lazy.difference' for the second map's leaves), or lookup (in +-- 'Data.IntMap.Lazy.intersection'). These helpers match the general structure +-- of normal single-key operations. However, there still need to be two +-- variants for which map's leaves are being inserted into which other map. +-- +-- If the 'Tip' corresponds to the bound that was already handled earlier in +-- the merge process, however, the recursion can immediately end. +-- +-- == Deletion and Choice of Intermediate +-- +-- Each step in the merge process as described so far processes a single key +-- (the merged bound for the node), recursion on the left branch (or a left +-- branch taken from one map), and recursion on the right brnach (or a right +-- branch taken from one map). This naturally corresponds to the arguments of +-- 'Bin' (a key/value pair and two 'Node's). However, that only matches up +-- when all of the keys are preserved, as in 'Data.IntMap.Lazy.union'. If the +-- merged bound is instead deleted, then it needs to be replaced with a bound +-- pulled from one of the two recursive cases. +-- +-- If the helper functions return 'Node's, extracting bounds from the recursive +-- cases using 'deleteMinL' and 'deleteMaxR' (or a wrapper function like +-- 'extractBinL' or 'extractBinR') is an expensive operation. After traversing +-- and reconstructing the subtrees in the merging provcess, they need to be +-- re-traversed and reconstructed again to pull out one of the entries. A more +-- efficient option would be to fuse the traversals, returning an 'IntMap_' +-- the recursive case. Since an 'IntMap_' already contains its own bound pulled +-- out, we can just put the pieces together in constant time. +-- +-- However, 'IntMap_' isn't a universally better choice for intermediate type. +-- When the outside key is kept, the 'IntMap_' needs to be converted back into +-- a 'Node', causing the same problem, but with insertion instead of deletion. +-- Both intermediate types are appropriate, often in the same merge operation, +-- depending only on whether the external bound has been kept or deleted. +-- 'Data.IntMap.Lazy.union' always keeps its keys so uses 'Node's everywhere. +-- 'Data.IntMap.Lazy.intersection' keeps its keys when they match, so +-- @go{L,R}Fused@ return 'Node's, but mismatched keys are dropped, so +-- @go{L,R}{1,2}@ return 'IntMap_'s. A function like +-- 'Data.IntMap.Lazy.differenceWith' needs two variants of @go{L,R}Fused@ +-- since different matched keys are kept and deleted in the same merge +-- operation. +-- +-- 'mergeA' is a particularly tricky function with regards to this decision. +-- Unlike the specialized merge operations, we don't statically know which keys +-- will be deleted, and unlike 'Data.IntMap.Lazy.differenceWith', we don't even +-- have that information at runtime: merge tactics return @f ('Maybe' a)@ +-- where @f@ is an arbitrary 'Applicative', and we can't choose what actions to +-- combine with that value based on the value itself. (As an aside, a 'Monad' +-- bound would not help either despite providing exactly that kind of choice, +-- as the actions associated with maximum bounds need to be sequenced after the +-- actions of the recursive cases.) Therefore, we need to choose a +-- representation with a good average case. If keys are being deleted, then the +-- maps returned by the recursive cases will be smaller, making the overhead of +-- pulling new bounds out also smaller. Therefore, universally using 'Node's +-- has a smaller overhead than universally using 'IntMap_'s and so is a better +-- choice. +-- +-- TODO: Theoretically, the recursive case could return both a 'Node' and an +-- 'IntMap_', stored in a lazy pair. This would allow selecting which to use +-- and only evaluating the one that is needed. It may also just result in +-- excessive work. +----------------------------------------------------------------------------- + +module Data.IntMap.Merge.Internal ( + -- * Types and inspection of tactics + WhenMissing(..) + , WhenMatched(..) + , SimpleWhenMissing + , SimpleWhenMatched +#if !MIN_VERSION_base(4,8,0) + , Identity(..) +#endif + , runWhenMissing + , runWhenMatched + , runWhenMissingAll + + -- * General combining functions + , merge + , mergeA + + -- * @WhenMissing@ tactics + , dropMissing + , preserveMissing + , filterMissing + , filterAMissing + , traverseMaybeMissingUKeyLazy +) where + +import Prelude hiding (min, max) + +import Data.IntMap.Internal + +import Control.Applicative (liftA2, liftA3) +import Control.Monad ((<=<)) +import qualified Control.Category as Category (Category(..)) +#if MIN_VERSION_base (4,8,0) +import Data.Functor.Identity (Identity, runIdentity) +#else +import Data.Functor ((<$)) +import Control.Applicative (Applicative(..), (<$>)) + +#if __GLASGOW_HASKELL__ >= 708 +import Data.Coerce +#endif + +-- | The identity type. +newtype Identity a = Identity { runIdentity :: a } + +#if __GLASGOW_HASKELL__ >= 708 +instance Functor Identity where + fmap = coerce +instance Applicative Identity where + (<*>) = coerce + pure = Identity +#else +instance Functor Identity where + fmap f (Identity a) = Identity (f a) +instance Applicative Identity where + Identity f <*> Identity x = Identity (f x) + pure = Identity +#endif +#endif + + +-- | A tactic for dealing with keys present in one map but not the other in +-- 'merge' or 'mergeA'. +-- +-- A tactic of type @ WhenMissing f a c @ is an abstract representation +-- of a function of type @ Key -> a -> f (Maybe c) @. +-- +-- @since 0.5.9 +data WhenMissing f a b = WhenMissing { + -- TODO: Instead of unconditionally putting the 'Maybe' inside the @f@, + -- provide options for the deletion vs. not choice to be decided purely, + -- before Applicative effects. Most merge tactics (all but 'filterAMissing' + -- and 'Data.IntMap.Merge.Lazy.traverseMaybeMissing', it seems) can support + -- this more precise representation, and it would allow 'mergeA' to make an + -- optimized choice between 'Node' and 'IntMap_' intermediates. By changing + -- its arguments into that form, 'merge' might even be implementable as + -- efficiently as possible in terms of 'mergeA'. + missingSingle :: UKey -> a -> f (Maybe b), + missingLeft :: Node L a -> f (Node L b), + missingRight :: Node R a -> f (Node R b), + missingAllL :: IntMap_ L a -> f (IntMap_ L b) +} + +-- | @since 0.5.9 +instance (Applicative f, Monad f) => Category.Category (WhenMissing f) where + {-# INLINE id #-} + id = preserveMissing + + {-# INLINE (.) #-} + missF . missG = WhenMissing + (\k a -> do + mb <- missingSingle missG k a + case mb of + Nothing -> return Nothing + Just b -> missingSingle missF k b) + (missingLeft missF <=< missingLeft missG) + (missingRight missF <=< missingRight missG) + (missingAllL missF <=< missingAllL missG) + +-- | @since 0.5.9 +instance (Applicative f, Monad f) => Functor (WhenMissing f a) where + {-# INLINE fmap #-} + fmap f miss = WhenMissing + (\k a -> fmap f <$> missingSingle miss k a) + (\l -> fmap f <$> missingLeft miss l) + (\r -> fmap f <$> missingRight miss r) + (\m -> fmap f <$> missingAllL miss m) + +-- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ +-- +-- @since 0.5.9 +instance (Applicative f, Monad f) => Applicative (WhenMissing f a) where + {-# INLINE pure #-} + pure x = WhenMissing + (\_ _ -> pure (Just x)) + (\l -> pure (x <$ l)) + (\r -> pure (x <$ r)) + (\m -> pure (x <$ m)) + + {-# INLINE (<*>) #-} + missF <*> missX = traverseMaybeMissingUKeyLazy $ \k a -> + liftA2 (<*>) (missingSingle missF k a) (missingSingle missX k a) + +-- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ +-- +-- @since 0.5.9 +instance (Applicative f, Monad f) => Monad (WhenMissing f a) where +#if !MIN_VERSION_base(4,8,0) + {-# INLINE return #-} + return = pure +#endif + + {-# INLINE (>>=) #-} + missM >>= f = traverseMaybeMissingUKeyLazy $ \k a -> do + mb <- missingSingle missM k a + case mb of + Nothing -> return Nothing + Just b -> missingSingle (f b) k a + +-- | The inverse of 'missingSingle'. Is in the @Internal@ module for type class +-- instances. +{-# INLINE traverseMaybeMissingUKeyLazy #-} +traverseMaybeMissingUKeyLazy :: Applicative f => (UKey -> a -> f (Maybe b)) -> WhenMissing f a b +traverseMaybeMissingUKeyLazy f = WhenMissing + { missingAllL = start + , missingLeft = goL + , missingRight = goR + , missingSingle = f } + where + start Empty = pure Empty + start (NonEmpty min minV root) = liftA2 (maybe nodeToMapL (NonEmpty min)) (f (boundUKey min) minV) (goL root) + + goL Tip = pure Tip + goL (Bin max maxV l r) = liftA3 (\l' r' maxV' -> maybe extractBinL (Bin max) maxV' l' r') (goL l) (goR r) (f (boundUKey max) maxV) + + goR Tip = pure Tip + goR (Bin min minV l r) = liftA3 (\minV' l' r' -> maybe extractBinR (Bin min) minV' l' r') (f (boundUKey min) minV) (goL l) (goR r) + +-- | A tactic for dealing with keys present in one map but not the other in +-- 'merge'. +-- +-- A tactic of type @ SimpleWhenMissing a c @ is an abstract representation +-- of a function of type @ Key -> a -> Maybe c @. +-- +-- @since 0.5.9 +type SimpleWhenMissing = WhenMissing Identity + +-- | Along with 'traverseMaybeMissing', witnesses the isomorphism +-- between @WhenMissing f x y@ and @Key -> x -> f (Maybe y)@. +-- +-- @since 0.5.9 +{-# INLINE runWhenMissing #-} +runWhenMissing :: WhenMissing f a b -> Key -> a -> f (Maybe b) +runWhenMissing miss k = missingSingle miss (unbox k) + +-- | Apply a missing tactic to an entire map. +-- +-- prop> runWhenMissingAll miss m = merge miss dropMissing dropMatched m empty +{-# INLINE runWhenMissingAll #-} +runWhenMissingAll :: Applicative f => WhenMissing f a b -> IntMap a -> f (IntMap b) +runWhenMissingAll miss (IntMap m) = IntMap <$> missingAllL miss m + +-- | Drop all the entries whose keys are missing from the other +-- map. +-- +-- @ +-- dropMissing :: SimpleWhenMissing a b +-- @ +-- +-- prop> dropMissing = mapMaybeMissing (\_ _ -> Nothing) +-- +-- but @dropMissing@ is much faster. +-- +-- @since 0.5.9 +{-# INLINE dropMissing #-} +dropMissing :: Applicative f => WhenMissing f a b +dropMissing = WhenMissing (\_ _ -> pure Nothing) (const (pure Tip)) (const (pure Tip)) (const (pure Empty)) + +-- | Preserve, unchanged, the entries whose keys are missing from +-- the other map. +-- +-- @ +-- preserveMissing :: SimpleWhenMissing a a +-- @ +-- +-- prop> preserveMissing = Merge.Lazy.mapMaybeMissing (\_ x -> Just x) +-- +-- but @preserveMissing@ is much faster. +-- +-- @since 0.5.9 +{-# INLINE preserveMissing #-} +preserveMissing :: Applicative f => WhenMissing f a a +preserveMissing = WhenMissing (\_ v -> pure (Just v)) pure pure pure + +-- | Filter the entries whose keys are missing from the other map. +-- +-- @ +-- filterMissing :: (Key -> x -> Bool) -> SimpleWhenMissing a a +-- @ +-- +-- 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 +{-# INLINE filterMissing #-} +filterMissing :: Applicative f => (Key -> a -> Bool) -> WhenMissing f a a +filterMissing p = filterMissingUKey (\k a -> p (box k) a) + +-- | Filter the entries whose keys are missing from the other map with a +-- predicate taking unboxed keys. Identical in functionality to +-- 'filterMissing'. +filterMissingUKey :: Applicative f => (UKey -> a -> Bool) -> WhenMissing f a a +filterMissingUKey p = WhenMissing (\k v -> pure (if p k v then Just v else Nothing)) (pure . goLKeep) (pure . goRKeep) (pure . start) where + start Empty = Empty + start (NonEmpty min minV root) + | p (boundUKey min) minV = NonEmpty min minV (goLKeep root) + | otherwise = goL root + + goLKeep Tip = Tip + goLKeep (Bin max maxV l r) + | p (boundUKey max) maxV = Bin max maxV (goLKeep l) (goRKeep r) + | otherwise = binNodeMapL (goLKeep l) (goR r) + + goRKeep Tip = Tip + goRKeep (Bin min minV l r) + | p (boundUKey min) minV = Bin min minV (goLKeep l) (goRKeep r) + | otherwise = binMapNodeR (goL l) (goRKeep r) + + goL Tip = Empty + goL (Bin max maxV l r) + | p (boundUKey max) maxV = binL (goL l) (NonEmpty max maxV (goRKeep r)) + | otherwise = binL (goL l) (goR r) + + goR Tip = Empty + goR (Bin min minV l r) + | p (boundUKey min) minV = binR (NonEmpty min minV (goLKeep l)) (goR r) + | otherwise = binR (goL l) (goR r) + +-- | 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 +{-# INLINE filterAMissing #-} +-- TODO: Use pointer equality to speed this up? +filterAMissing :: Applicative f => (Key -> a -> f Bool) -> WhenMissing f a a +filterAMissing p = filterAMissingUKey (\k a -> p (box k) a) + +-- | Filter the entries whose keys are missing from the other map using some +-- 'Applicative' action that takes an unboxed key. Identical in functionality +-- to 'filterAMissing'. +filterAMissingUKey :: Applicative f => (UKey -> a -> f Bool) -> WhenMissing f a a +filterAMissingUKey f = WhenMissing + { missingAllL = start + , missingLeft = goL + , missingRight = goR + , missingSingle = \k v -> fmap (\keep -> if keep then Just v else Nothing) (f k v) } + where + start Empty = pure Empty + start (NonEmpty min minV root) = liftA2 (\keepV root' -> if keepV then NonEmpty min minV root' else nodeToMapL root') (f (boundUKey min) minV) (goL root) + + goL Tip = pure Tip + goL (Bin max maxV l r) = liftA3 (\l' r' keepMax -> if keepMax then Bin max maxV l' r' else extractBinL l' r') (goL l) (goR r) (f (boundUKey max) maxV) + + goR Tip = pure Tip + goR (Bin min minV l r) = liftA3 (\keepMin l' r' -> if keepMin then Bin min minV l' r' else extractBinR l' r') (f (boundUKey min) minV) (goL l) (goR r) + +-- | A tactic for dealing with keys present in both +-- maps in 'merge' or 'mergeA'. +-- +-- A tactic of type @ WhenMatched f a b c @ is an abstract representation +-- of a function of type @ Key -> a -> b -> f (Maybe c) @. +-- +-- @since 0.5.9 +newtype WhenMatched f a b c = WhenMatched { + matchedSingle :: UKey -> a -> b -> f (Maybe c) +} + +-- | @since 0.5.9 +instance (Applicative f, Monad f) => Category.Category (WhenMatched f a) where + -- TODO: Expose this and it symmetric pair as @firstMatched@ and + -- @secondMatched@? + {-# INLINE id #-} + id = WhenMatched (\_ _ b -> pure (Just b)) + + {-# INLINE (.) #-} + matchF . matchG = WhenMatched $ \k a b -> do + mc <- matchedSingle matchG k a b + case mc of + Nothing -> return Nothing + Just c' -> matchedSingle matchF k a c' + +-- | @since 0.5.9 +instance Functor f => Functor (WhenMatched f a b) where + {-# INLINE fmap #-} + fmap f match = WhenMatched (\k a b -> fmap f <$> matchedSingle match k a b) + +-- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ +-- +-- @since 0.5.9 +instance (Applicative f, Monad f) => Applicative (WhenMatched f a b) where + {-# INLINE pure #-} + pure x = WhenMatched (\_ _ _ -> pure (Just x)) + + {-# INLINE (<*>) #-} + matchF <*> matchX = WhenMatched $ \k a b -> + liftA2 (<*>) (matchedSingle matchF k a b) (matchedSingle matchX k a b) + +-- | Equivalent to @ReaderT Key (ReaderT x (ReaderT y (MaybeT f)))@ +-- +-- @since 0.5.9 +instance (Applicative f, Monad f) => Monad (WhenMatched f a b) where +#if !MIN_VERSION_base(4,8,0) + {-# INLINE return #-} + return = pure +#endif + + {-# INLINE (>>=) #-} + matchM >>= f = WhenMatched $ \k a b -> do + mc <- matchedSingle matchM k a b + case mc of + Nothing -> return Nothing + Just c -> matchedSingle (f c) k a b + +-- | A tactic for dealing with keys present in both maps in 'merge'. +-- +-- A tactic of type @ SimpleWhenMatched a b c @ is an abstract representation +-- of a function of type @ Key -> a -> b -> Maybe c @. +-- +-- @since 0.5.9 +type SimpleWhenMatched = WhenMatched Identity + +-- | Along with 'zipWithMaybeAMatched', witnesses the isomorphism +-- between @WhenMatched f x y z@ and @Key -> x -> y -> f (Maybe z)@. +-- +-- @since 0.5.9 +{-# INLINE runWhenMatched #-} +runWhenMatched :: WhenMatched f a b c -> Key -> a -> b -> f (Maybe c) +runWhenMatched match k = matchedSingle match (unbox k) + +-- | 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 preserveMissing dropMissing f +-- prop> symmetricDifference = merge preserveMissing preserveMissing (zipWithMaybeMatched (\_ _ _ -> Nothing)) +-- prop> mapEachPiece f g h = merge (mapMissing f) (mapMissing g) (zipWithMatched h) +-- +-- @since 0.5.9 +{-# INLINE merge #-} +-- TODO: Implementing 'merge' in terms of 'mergeA' leaves a lot to be desired +-- on the performance front. Because the choice of whether to delete a key is +-- hidden behind Applicative effects, 'mergeA' must conservatively return +-- 'Node's from every intermediate function, even when an 'IntMap_' would be +-- more efficient and appropriate. (It could return 'IntMap_'s everywhere +-- instead, but that would have the same problem in the opposite direction.) +-- 'merge' has no such limitation, since recursion schemes and decisions about +-- what types of structures to return can be made based on the results of pure +-- tactics. +merge :: SimpleWhenMissing a c -> SimpleWhenMissing b c -> SimpleWhenMatched a b c -> IntMap a -> IntMap b -> IntMap c +merge miss1 miss2 match = \m1 m2 -> runIdentity (mergeA miss1 miss2 match m1 m2) + +-- | 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 +{-# INLINE mergeA #-} +mergeA :: Applicative f => WhenMissing f a c -> WhenMissing f b c -> WhenMatched f a b c -> IntMap a -> IntMap b -> f (IntMap c) +mergeA miss1 miss2 match = start where + start (IntMap Empty) (IntMap Empty) = pure (IntMap Empty) + start (IntMap Empty) (IntMap !m2) = IntMap <$> missingAllL miss2 m2 + start (IntMap !m1) (IntMap Empty) = IntMap <$> missingAllL miss1 m1 + start (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = makeIntMapNE min1 (missingSingle miss1 (boundUKey min1) minV1) (goL2 minV2 min1 root1 min2 root2) + | min2 < min1 = makeIntMapNE min2 (missingSingle miss2 (boundUKey min2) minV2) (goL1 minV1 min1 root1 min2 root2) + | otherwise = makeIntMapNE min1 (matchedSingle match (boundUKey min1) minV1 minV2) (goLFused min1 root1 root2) + + goL1 minV1 !min1 !n1 !min2 Tip = mapToNodeL min2 <$> missingAllL miss1 (NonEmpty min1 minV1 n1) + goL1 minV1 !min1 !n1 !min2 n2@(Bin max2 _ _ _) | boundsDisjoint min1 max2 = liftA2 (maybeUnionDisjointL min2) (missingLeft miss2 n2) (missingAllL miss1 (NonEmpty min1 minV1 n1)) + goL1 minV1 !min1 Tip !min2 !n2 = goInsertL1 (boundKey min1) minV1 (xor (boundKey min1) min2) min2 n2 + goL1 minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> liftA2 binNodeMapL (goL1 minV1 min1 n1 min2 l2) (missingAllR miss2 (NonEmpty max2 maxV2 r2)) + | max1 > max2 -> makeBinL max1 (missingSingle miss1 (boundUKey max1) maxV1) (missingLeft miss2 l2) (goR2 maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> makeBinL max2 (missingSingle miss2 (boundUKey max2) maxV2) (missingLeft miss2 l2) (goR1 maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2) + | otherwise -> makeBinL max1 (matchedSingle match (boundUKey max1) maxV1 maxV2) (missingLeft miss2 l2) (goRFused max1 (Bin min1 minV1 l1 r1) r2) + EQ | max1 > max2 -> makeBinL max1 (missingSingle miss1 (boundUKey max1) maxV1) (goL1 minV1 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | max1 < max2 -> makeBinL max2 (missingSingle miss2 (boundUKey max2) maxV2) (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> makeBinL max1 (matchedSingle match (boundUKey max1) maxV1 maxV2) (goL1 minV1 min1 l1 min2 l2) (goRFused max1 r1 r2) + GT -> liftA2 binNodeMapL (goL1 minV1 min1 l1 min2 n2) (missingAllR miss1 (NonEmpty max1 maxV1 r1)) + + goL2 minV2 !min1 Tip !min2 !n2 = mapToNodeL min1 <$> missingAllL miss2 (NonEmpty min2 minV2 n2) + goL2 minV2 !min1 n1@(Bin max1 _ _ _) !min2 !n2 | boundsDisjoint min2 max1 = liftA2 (maybeUnionDisjointL min1) (missingLeft miss1 n1) (missingAllL miss2 (NonEmpty min2 minV2 n2)) + goL2 minV2 !min1 !n1 !min2 Tip = goInsertL2 (boundKey min2) minV2 (xor (boundKey min2) min1) min1 n1 + goL2 minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> liftA2 binNodeMapL (goL2 minV2 min1 l1 min2 n2) (missingAllR miss1 (NonEmpty max1 maxV1 r1)) + | max1 > max2 -> makeBinL max1 (missingSingle miss1 (boundUKey max1) maxV1) (missingLeft miss1 l1) (goR2 maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | max1 < max2 -> makeBinL max2 (missingSingle miss2 (boundUKey max2) maxV2) (missingLeft miss1 l1) (goR1 maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | otherwise -> makeBinL max1 (matchedSingle match (boundUKey max1) maxV1 maxV2) (missingLeft miss1 l1) (goRFused max1 r1 (Bin min2 minV2 l2 r2)) + EQ | max1 > max2 -> makeBinL max1 (missingSingle miss1 (boundUKey max1) maxV1) (goL2 minV2 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | max1 < max2 -> makeBinL max2 (missingSingle miss2 (boundUKey max2) maxV2) (goL2 minV2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> makeBinL max1 (matchedSingle match (boundUKey max1) maxV1 maxV2) (goL2 minV2 min1 l1 min2 l2) (goRFused max1 r1 r2) + LT -> liftA2 binNodeMapL (goL2 minV2 min1 n1 min2 l2) (missingAllR miss2 (NonEmpty max2 maxV2 r2)) + + goLFused !_ Tip !n2 = missingLeft miss2 n2 + goLFused !_ !n1 Tip = missingLeft miss1 n1 + goLFused !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> liftA2 binNodeMapL (goLFused min n1 l2) (missingAllR miss2 (NonEmpty max2 maxV2 r2)) + EQ | max1 > max2 -> makeBinL max1 (missingSingle miss1 (boundUKey max1) maxV1) (goLFused min l1 l2) (goR2 maxV2 max1 r1 max2 r2) + | max1 < max2 -> makeBinL max2 (missingSingle miss2 (boundUKey max2) maxV2) (goLFused min l1 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> makeBinL max1 (matchedSingle match (boundUKey max1) maxV1 maxV2) (goLFused min l1 l2) (goRFused max1 r1 r2) + GT -> liftA2 binNodeMapL (goLFused min l1 n2) (missingAllR miss1 (NonEmpty max1 maxV1 r1)) + + goR1 maxV1 !max1 !n1 !max2 Tip = mapToNodeR max2 <$> missingAllR miss1 (NonEmpty max1 maxV1 n1) + goR1 maxV1 !max1 !n1 !max2 n2@(Bin min2 _ _ _) | boundsDisjoint min2 max1 = liftA2 (maybeUnionDisjointR max2) (missingAllR miss1 (NonEmpty max1 maxV1 n1)) (missingRight miss2 n2) + goR1 maxV1 !max1 Tip !max2 !n2 = goInsertR1 (boundKey max1) maxV1 (xor (boundKey max1) max2) max2 n2 + goR1 maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> liftA2 binMapNodeR (missingAllL miss2 (NonEmpty min2 minV2 l2)) (goR1 maxV1 max1 n1 max2 r2) + | min1 < min2 -> makeBinR min1 (missingSingle miss1 (boundUKey min1) minV1) (goL2 minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) (missingRight miss2 r2) + | min1 > min2 -> makeBinR min2 (missingSingle miss2 (boundUKey min2) minV2) (goL1 minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2) (missingRight miss2 r2) + | otherwise -> makeBinR min1 (matchedSingle match (boundUKey min1) minV1 minV2) (goLFused min1 (Bin max1 maxV1 l1 r1) l2) (missingRight miss2 r2) + EQ | min1 < min2 -> makeBinR min1 (missingSingle miss1 (boundUKey min1) minV1) (goL2 minV2 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | min1 > min2 -> makeBinR min2 (missingSingle miss2 (boundUKey min2) minV2) (goL1 minV1 min1 l1 min2 l2) (goR1 maxV1 max1 r1 max2 r2) + | otherwise -> makeBinR min1 (matchedSingle match (boundUKey min1) minV1 minV2) (goLFused min1 l1 l2) (goR1 maxV1 max1 r1 max2 r2) + GT -> liftA2 binMapNodeR (missingAllL miss1 (NonEmpty min1 minV1 l1)) (goR1 maxV1 max1 r1 max2 n2) + + goR2 maxV2 !max1 Tip !max2 !n2 = mapToNodeR max1 <$> missingAllR miss2 (NonEmpty max2 maxV2 n2) + goR2 maxV2 !max1 n1@(Bin min1 _ _ _) !max2 !n2 | boundsDisjoint min1 max2 = liftA2 (maybeUnionDisjointR max1) (missingAllR miss2 (NonEmpty max2 maxV2 n2)) (missingRight miss1 n1) + goR2 maxV2 !max1 !n1 !max2 Tip = goInsertR2 (boundKey max2) maxV2 (xor (boundKey max1) max2) max1 n1 + goR2 maxV2 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> liftA2 binMapNodeR (missingAllL miss1 (NonEmpty min1 minV1 l1)) (goR2 maxV2 max1 r1 max2 n2) + | min1 < min2 -> makeBinR min1 (missingSingle miss1 (boundUKey min1) minV1) (goL2 minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2)) (missingRight miss1 r1) + | min1 > min2 -> makeBinR min2 (missingSingle miss2 (boundUKey min2) minV2) (goL1 minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2)) (missingRight miss1 r1) + | otherwise -> makeBinR min1 (matchedSingle match (boundUKey min1) minV1 minV2) (goLFused min1 l1 (Bin max2 maxV2 l2 r2)) (missingRight miss1 r1) + EQ | min1 < min2 -> makeBinR min1 (missingSingle miss1 (boundUKey min1) minV1) (goL2 minV2 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | min1 > min2 -> makeBinR min2 (missingSingle miss2 (boundUKey min2) minV2) (goL1 minV1 min1 l1 min2 l2) (goR2 maxV2 max1 r1 max2 r2) + | otherwise -> makeBinR min1 (matchedSingle match (boundUKey min1) minV1 minV2) (goLFused min1 l1 l2) (goR2 maxV2 max1 r1 max2 r2) + LT -> liftA2 binMapNodeR (missingAllL miss2 (NonEmpty min2 minV2 l2)) (goR2 maxV2 max1 n1 max2 r2) + + goRFused !_ Tip !n2 = missingRight miss2 n2 + goRFused !_ !n1 Tip = missingRight miss1 n1 + goRFused !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> liftA2 binMapNodeR (missingAllL miss2 (NonEmpty min2 minV2 l2)) (goRFused max n1 r2) + EQ | min1 < min2 -> makeBinR min1 (missingSingle miss1 (boundUKey min1) minV1) (goL2 minV2 min1 l1 min2 l2) (goRFused max r1 r2) + | min1 > min2 -> makeBinR min2 (missingSingle miss2 (boundUKey min2) minV2) (goL1 minV1 min1 l1 min2 l2) (goRFused max r1 r2) + | otherwise -> makeBinR min1 (matchedSingle match (boundUKey min1) minV1 minV2) (goLFused min1 l1 l2) (goRFused max r1 r2) + GT -> liftA2 binMapNodeR (missingAllL miss1 (NonEmpty min1 minV1 l1)) (goRFused max r1 n2) + + goInsertL1 !k v !_ !_ Tip = makeSingleton k (missingSingle miss1 (unbox k) v) + goInsertL1 !k v !xorCache !min n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> liftA2 binNodeMapL (goInsertL1 k v xorCache min l) (missingAllR miss2 (NonEmpty max maxV r)) + | otherwise -> makeBinL max (missingSingle miss2 (boundUKey max) maxV) (missingLeft miss2 l) (goInsertR1 k v xorCacheMax max r) + OutOfBound -> addMaxL min (Bound k) (missingSingle miss1 (unbox k) v) (missingLeft miss2 n) + Matched -> makeBinL max (matchedSingle match (unbox k) v maxV) (missingLeft miss2 l) (missingRight miss2 r) + where xorCacheMax = xor k max + + goInsertL2 !k v !_ !_ Tip = makeSingleton k (missingSingle miss2 (unbox k) v) + goInsertL2 !k v !xorCache !min n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> liftA2 binNodeMapL (goInsertL2 k v xorCache min l) (missingAllR miss1 (NonEmpty max maxV r)) + | otherwise -> makeBinL max (missingSingle miss1 (boundUKey max) maxV) (missingLeft miss1 l) (goInsertR2 k v xorCacheMax max r) + OutOfBound -> addMaxL min (Bound k) (missingSingle miss2 (unbox k) v) (missingLeft miss1 n) + Matched -> makeBinL max (matchedSingle match (unbox k) maxV v) (missingLeft miss1 l) (missingRight miss1 r) + where xorCacheMax = xor k max + + goInsertR1 !k v !_ !_ Tip = makeSingleton k (missingSingle miss1 (unbox k) v) + goInsertR1 !k v !xorCache !max n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> liftA2 binMapNodeR (missingAllL miss2 (NonEmpty min minV l)) (goInsertR1 k v xorCache max r) + | otherwise -> makeBinR min (missingSingle miss2 (boundUKey min) minV) (goInsertL1 k v xorCacheMin min l) (missingRight miss2 r) + OutOfBound -> addMinR max (Bound k) (missingSingle miss1 (unbox k) v) (missingRight miss2 n) + Matched -> makeBinR min (matchedSingle match (unbox k) v minV) (missingLeft miss2 l) (missingRight miss2 r) + where xorCacheMin = xor k min + + goInsertR2 !k v !_ !_ Tip = makeSingleton k (missingSingle miss2 (unbox k) v) + goInsertR2 !k v !xorCache !max n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> liftA2 binMapNodeR (missingAllL miss1 (NonEmpty min minV l)) (goInsertR2 k v xorCache max r) + | otherwise -> makeBinR min (missingSingle miss1 (boundUKey min) minV) (goInsertL2 k v xorCacheMin min l) (missingRight miss1 r) + OutOfBound -> addMinR max (Bound k) (missingSingle miss2 (unbox k) v) (missingRight miss1 n) + Matched -> makeBinR min (matchedSingle match (unbox k) minV v) (missingLeft miss1 l) (missingRight miss1 r) + where xorCacheMin = xor k min + + missingAllR whenMiss = fmap l2rMap . missingAllL whenMiss . r2lMap + +{-# INLINE makeSingleton #-} +makeSingleton :: Functor f => Key -> f (Maybe v) -> f (Node d v) +makeSingleton !k v = make <$> v where + make Nothing = Tip + make (Just v') = Bin (Bound k) v' Tip Tip + +{-# INLINE makeBinL #-} +makeBinL :: Applicative f => Bound R -> f (Maybe v) -> f (Node L v) -> f (Node R v) -> f (Node L v) +makeBinL !max maxV l r = liftA3 make l r maxV where + make l' r' maxV' = maybe extractBinL (Bin max) maxV' l' r' + +{-# INLINE makeBinR #-} +makeBinR :: Applicative f => Bound L -> f (Maybe v) -> f (Node L v) -> f (Node R v) -> f (Node R v) +makeBinR !min minV l r = liftA3 make minV l r where + make minV' l' r' = maybe extractBinR (Bin min) minV' l' r' + +{-# INLINE makeIntMapNE #-} +makeIntMapNE :: Applicative f => Bound L -> f (Maybe v) -> f (Node L v) -> f (IntMap v) +makeIntMapNE !min minV root = liftA2 make minV root where + make Nothing root' = IntMap (nodeToMapL root') + make (Just minV') root' = IntMap (NonEmpty min minV' root') + +mapToNodeL :: Bound L -> IntMap_ L v -> Node L v +mapToNodeL !_ Empty = Tip +mapToNodeL !externalMin (NonEmpty min minV root) = insertMinL (xor (boundKey min) externalMin) min minV root + +mapToNodeR :: Bound R -> IntMap_ R v -> Node R v +mapToNodeR !_ Empty = Tip +mapToNodeR !externalMax (NonEmpty max maxV root) = insertMaxR (xor (boundKey max) externalMax) max maxV root + +{-# INLINE addMaxL #-} +addMaxL :: Applicative f => Bound L -> Bound R -> f (Maybe v) -> f (Node L v) -> f (Node L v) +addMaxL !min !k v n = liftA2 add n v where + add n' Nothing = n' + add Tip (Just v') = Bin k v' Tip Tip + add n'@(Bin max maxV l r) (Just v') + | xor (boundKey max) min < xorCacheMax = Bin k v' n' Tip + | otherwise = Bin k v' l (insertMaxR xorCacheMax max maxV r) + where + xorCacheMax = xor (boundKey max) k + +{-# INLINE addMinR #-} +addMinR :: Applicative f => Bound R -> Bound L -> f (Maybe v) -> f (Node R v) -> f (Node R v) +addMinR !max !k v n = liftA2 add v n where + add Nothing n' = n' + add (Just v') Tip = Bin k v' Tip Tip + add (Just v') n'@(Bin min minV l r) + | xor (boundKey min) max < xorCacheMin = Bin k v' Tip n' + | otherwise = Bin k v' (insertMinL xorCacheMin min minV l) r + where + xorCacheMin = xor (boundKey min) k + +maybeUnionDisjointL :: Bound L -> Node L v -> IntMap_ L v -> Node L v +maybeUnionDisjointL !min1 Tip !m2 = mapToNodeL min1 m2 +maybeUnionDisjointL !_ !n1 Empty = n1 +maybeUnionDisjointL !min1 !n1 (NonEmpty min2 minV2 root2) = unionDisjointL minV2 min1 n1 min2 root2 + +maybeUnionDisjointR :: Bound R -> IntMap_ R v -> Node R v -> Node R v +maybeUnionDisjointR !max2 !m1 Tip = mapToNodeR max2 m1 +maybeUnionDisjointR !_ Empty !n2 = n2 +maybeUnionDisjointR !max2 (NonEmpty max1 maxV1 root1) !n2 = unionDisjointR maxV1 max1 root1 max2 n2 diff --git a/containers/src/Data/IntMap/Merge/Lazy.hs b/containers/src/Data/IntMap/Merge/Lazy.hs index c24d0e45f..99dc88649 100644 --- a/containers/src/Data/IntMap/Merge/Lazy.hs +++ b/containers/src/Data/IntMap/Merge/Lazy.hs @@ -1,40 +1,27 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE BangPatterns #-} -#if __GLASGOW_HASKELL__ -{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} -#endif +{-# LANGUAGE CPP, BangPatterns #-} #if !defined(TESTING) && defined(__GLASGOW_HASKELL__) {-# LANGUAGE Safe #-} #endif -#if __GLASGOW_HASKELL__ >= 708 -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE TypeFamilies #-} -#define USE_MAGIC_PROXY 1 -#endif - -#if USE_MAGIC_PROXY -{-# LANGUAGE MagicHash #-} -#endif - -#include "containers.h" ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Merge.Lazy --- Copyright : (c) wren romano 2016 +-- Copyright : Documentation & Interface (c) wren romano 2016 +-- Documentation & Implementation (c) Jonathan "gereeter" S. 2020 -- License : BSD-style -- Maintainer : libraries@haskell.org +-- Stability : provisional -- Portability : portable -- -- This module defines an API for writing functions that merge two -- maps. The key functions are 'merge' and 'mergeA'. --- Each of these can be used with several different \"merge tactics\". +-- 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.Strict.mapMissing' --- from "Data.Map.Merge.Strict" then the results will be forced before --- they are inserted. If you use 'Data.Map.Merge.Lazy.mapMissing' from +-- determines strictness. If you use 'Data.IntMap.Merge.Strict.mapMissing' +-- from "Data.IntMap.Merge.Strict" then the results will be forced before +-- they are inserted. If you use 'Data.IntMap.Merge.Lazy.mapMissing' from -- this module then they will not. -- -- == Efficiency note @@ -45,7 +32,6 @@ -- for 'WhenMatched' tactics should not pose any major efficiency problems. -- -- @since 0.5.9 - module Data.IntMap.Merge.Lazy ( -- ** Simple merge tactic types SimpleWhenMissing @@ -59,10 +45,10 @@ module Data.IntMap.Merge.Lazy ( , zipWithMatched -- *** @WhenMissing@ tactics - , mapMaybeMissing , dropMissing , preserveMissing , mapMissing + , mapMaybeMissing , filterMissing -- ** Applicative merge tactic types @@ -87,18 +73,225 @@ module Data.IntMap.Merge.Lazy ( , traverseMissing , filterAMissing - -- *** Covariant maps for tactics + -- ** Covariant maps for tactics , mapWhenMissing , mapWhenMatched - -- *** Contravariant maps for tactics + -- ** Contravariant maps for tactics , lmapWhenMissing , contramapFirstWhenMatched , contramapSecondWhenMatched - -- *** Miscellaneous tactic functions + -- ** Miscellaneous functions on tactics , runWhenMatched , runWhenMissing - ) where +) where + +#if !MIN_VERSION_base(4,8,0) +import Control.Applicative (Applicative(..), (<$>)) +#endif +import Control.Applicative (liftA2, liftA3) + +import Prelude hiding (min, max) import Data.IntMap.Internal +import Data.IntMap.Merge.Internal + +-- | Map over the entries whose keys are missing from the other map. +-- +-- @ +-- mapMissing :: (Key -> a -> b) -> SimpleWhenMissing a b +-- @ +-- +-- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- +-- but @mapMissing@ is somewhat faster. +-- +-- @since 0.5.9 +{-# INLINE mapMissing #-} +mapMissing :: Applicative f => (Key -> a -> b) -> WhenMissing f a b +mapMissing f = mapMissingUKey (\k v -> f (box k) v) + +-- | Map over the entries whose keys are missing from the other map with a +-- function that takes an unboxed key. Identical in functionality to +-- 'mapMissing'. +mapMissingUKey :: Applicative f => (UKey -> a -> b) -> WhenMissing f a b +mapMissingUKey g = WhenMissing (\k v -> pure (Just (g k v))) (pure . go g) (pure . go g) (pure . start g) where + start _ Empty = Empty + start f (NonEmpty min minV root) = NonEmpty min (f (boundUKey min) minV) (go f root) + + go :: (UKey -> a -> b) -> Node t a -> Node t b + go _ Tip = Tip + go f (Bin k v l r) = Bin k (f (boundUKey k) v) (go f l) (go f r) + +-- | 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 -> a -> Maybe b) -> SimpleWhenMissing a b +-- @ +-- +-- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) +-- +-- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' operations. +-- +-- @since 0.5.9 +{-# INLINE mapMaybeMissing #-} +mapMaybeMissing :: Applicative f => (Key -> a -> Maybe b) -> WhenMissing f a b +mapMaybeMissing f = mapMaybeMissingUKey (\k a -> f (box k) a) + +-- | Map over the entries whose keys are missing from the other map using a +-- function taking an unboxed key, optionally removing some. Identical in +-- functionality to 'mapMaybeMissing'. +mapMaybeMissingUKey :: Applicative f => (UKey -> a -> Maybe b) -> WhenMissing f a b +mapMaybeMissingUKey f = WhenMissing (\k v -> pure (f k v)) (pure . goLKeep) (pure . goRKeep) (pure . start) where + start Empty = Empty + start (NonEmpty min minV root) = case f (boundUKey min) minV of + Just minV' -> NonEmpty min minV' (goLKeep root) + Nothing -> goL root + + goLKeep Tip = Tip + goLKeep (Bin max maxV l r) = case f (boundUKey max) maxV of + Just maxV' -> Bin max maxV' (goLKeep l) (goRKeep r) + Nothing -> binNodeMapL (goLKeep l) (goR r) + + goRKeep Tip = Tip + goRKeep (Bin min minV l r) = case f (boundUKey min) minV of + Just minV' -> Bin min minV' (goLKeep l) (goRKeep r) + Nothing -> binMapNodeR (goL l) (goRKeep r) + + goL Tip = Empty + goL (Bin max maxV l r) = case f (boundUKey max) maxV of + Just maxV' -> binL (goL l) (NonEmpty max maxV' (goRKeep r)) + Nothing -> binL (goL l) (goR r) + + goR Tip = Empty + goR (Bin min minV l r) = case f (boundUKey min) minV of + Just minV' -> binR (NonEmpty min minV' (goLKeep l)) (goR r) + Nothing -> binR (goL l) (goR r) + +-- | 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 -> a -> b -> Maybe c) +-- -> SimpleWhenMatched a b c +-- @ +-- +-- @since 0.5.9 +{-# INLINE zipWithMaybeMatched #-} +zipWithMaybeMatched :: Applicative f => (Key -> a -> b -> Maybe c) -> WhenMatched f a b c +zipWithMaybeMatched f = WhenMatched (\k a b -> pure (f (box k) a b)) + +-- | 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 -> a -> b -> c) +-- -> SimpleWhenMatched a b c +-- @ +-- +-- @since 0.5.9 +{-# INLINE zipWithMatched #-} +zipWithMatched :: Applicative f => (Key -> a -> b -> c) -> WhenMatched f a b c +zipWithMatched f = zipWithMaybeMatched (\k a b -> Just (f k a b)) + +-- | 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 +{-# INLINE zipWithMaybeAMatched #-} +zipWithMaybeAMatched + :: (Key -> a -> b -> f (Maybe c)) + -> WhenMatched f a b c +zipWithMaybeAMatched f = WhenMatched (\k a b -> f (box k) a b) + +-- | 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 +{-# INLINE zipWithAMatched #-} +zipWithAMatched + :: Applicative f + => (Key -> a -> b -> f c) + -> WhenMatched f a b c +zipWithAMatched f = zipWithMaybeAMatched (\k a b -> Just <$> f k a b) + +-- | 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 +{-# INLINE traverseMaybeMissing #-} +traverseMaybeMissing + :: Applicative f => (Key -> a -> f (Maybe b)) -> WhenMissing f a b +traverseMaybeMissing f = traverseMaybeMissingUKeyLazy (\k a -> f (box k) a) + +-- | Traverse over the entries whose keys are missing from the other +-- map. +-- +-- @since 0.5.9 +{-# INLINE traverseMissing #-} +traverseMissing + :: Applicative f => (Key -> a -> f b) -> WhenMissing f a b +traverseMissing f = WhenMissing + { missingAllL = start + , missingLeft = goL + , missingRight = goR + , missingSingle = \k v -> Just <$> f' k v } + where + f' k a = f (box k) a + + start Empty = pure Empty + start (NonEmpty min minV root) = liftA2 (NonEmpty min) (f' (boundUKey min) minV) (goL root) + + goL Tip = pure Tip + goL (Bin max maxV l r) = liftA3 (\l' r' maxV' -> Bin max maxV' l' r') (goL l) (goR r) (f' (boundUKey max) maxV) + + goR Tip = pure Tip + goR (Bin min minV l r) = liftA3 (\minV' l' r' -> Bin min minV' l' r') (f' (boundUKey min) minV) (goL l) (goR r) + +-- | Map covariantly over a @'WhenMissing' f x@. +-- +-- @since 0.5.9 +{-# INLINE mapWhenMissing #-} +mapWhenMissing :: (Applicative f, Monad f) => (a -> b) -> WhenMissing f x a -> WhenMissing f x b +mapWhenMissing = fmap + +-- | Map covariantly over a @'WhenMatched' f x y@. +-- +-- @since 0.5.9 +{-# INLINE mapWhenMatched #-} +mapWhenMatched :: Functor f => (a -> b) -> WhenMatched f x y a -> WhenMatched f x y b +mapWhenMatched = fmap + +-- | Map contravariantly over a @'WhenMissing' f _ x@. +-- +-- @since 0.5.9 +{-# INLINE lmapWhenMissing #-} +lmapWhenMissing :: (b -> a) -> WhenMissing f a x -> WhenMissing f b x +lmapWhenMissing f miss = WhenMissing + (\k b -> missingSingle miss k (f b)) + (\l -> missingLeft miss (fmap f l)) + (\r -> missingRight miss (fmap f r)) + (\m -> missingAllL miss (fmap f m)) + +-- | Map contravariantly over a @'WhenMatched' f _ y z@. +-- +-- @since 0.5.9 +{-# INLINE contramapFirstWhenMatched #-} +contramapFirstWhenMatched :: (b -> a) -> WhenMatched f a y z -> WhenMatched f b y z +contramapFirstWhenMatched f match = WhenMatched (\k b y -> matchedSingle match k (f b) y) + +-- | Map contravariantly over a @'WhenMatched' f x _ z@. +-- +-- @since 0.5.9 +{-# INLINE contramapSecondWhenMatched #-} +contramapSecondWhenMatched :: (b -> a) -> WhenMatched f x a z -> WhenMatched f x b z +contramapSecondWhenMatched f match = WhenMatched (\k x b -> matchedSingle match k x (f b)) diff --git a/containers/src/Data/IntMap/Merge/Strict.hs b/containers/src/Data/IntMap/Merge/Strict.hs index d21c4e1ca..211aaea2d 100644 --- a/containers/src/Data/IntMap/Merge/Strict.hs +++ b/containers/src/Data/IntMap/Merge/Strict.hs @@ -1,40 +1,27 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE BangPatterns #-} -#if __GLASGOW_HASKELL__ -{-# LANGUAGE DeriveDataTypeable, StandaloneDeriving #-} -#endif +{-# LANGUAGE CPP, BangPatterns #-} #if !defined(TESTING) && defined(__GLASGOW_HASKELL__) -{-# LANGUAGE Trustworthy #-} -#endif -#if __GLASGOW_HASKELL__ >= 708 -{-# LANGUAGE RoleAnnotations #-} -{-# LANGUAGE TypeFamilies #-} -#define USE_MAGIC_PROXY 1 -#endif - -#if USE_MAGIC_PROXY -{-# LANGUAGE MagicHash #-} +{-# LANGUAGE Safe #-} #endif -#include "containers.h" - ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Merge.Strict --- Copyright : (c) wren romano 2016 +-- Copyright : Documentation & Interface (c) wren romano 2016 +-- Documentation & Implementation (c) Jonathan "gereeter" S. 2020 -- License : BSD-style -- Maintainer : libraries@haskell.org +-- Stability : provisional -- Portability : portable -- -- This module defines an API for writing functions that merge two -- maps. The key functions are 'merge' and 'mergeA'. --- Each of these can be used with several different \"merge tactics\". +-- 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.Strict.mapMissing' +-- determines strictness. If you use 'Data.IntMap.Merge.Strict.mapMissing' -- from this module then the results will be forced before they are --- inserted. If you use 'Data.Map.Merge.Lazy.mapMissing' from +-- inserted. If you use 'Data.IntMap.Merge.Lazy.mapMissing' from -- "Data.Map.Merge.Lazy" then they will not. -- -- == Efficiency note @@ -43,8 +30,6 @@ -- 'WhenMissing' tactics are included because they are valid. However, they are -- inefficient in many cases and should usually be avoided. The instances -- for 'WhenMatched' tactics should not pose any major efficiency problems. --- --- @since 0.5.9 module Data.IntMap.Merge.Strict ( -- ** Simple merge tactic types @@ -59,10 +44,10 @@ module Data.IntMap.Merge.Strict ( , zipWithMatched -- *** @WhenMissing@ tactics - , mapMaybeMissing , dropMissing , preserveMissing , mapMissing + , mapMaybeMissing , filterMissing -- ** Applicative merge tactic types @@ -92,142 +77,215 @@ module Data.IntMap.Merge.Strict ( , mapWhenMatched -- ** Miscellaneous functions on tactics - , runWhenMatched , runWhenMissing - ) where +) where -import Data.IntMap.Internal - ( SimpleWhenMissing - , SimpleWhenMatched - , merge - , dropMissing - , preserveMissing - , filterMissing - , WhenMissing (..) - , WhenMatched (..) - , mergeA - , filterAMissing - , runWhenMatched - , runWhenMissing - ) -import Data.IntMap.Strict.Internal #if !MIN_VERSION_base(4,8,0) -import Control.Applicative (Applicative (..), (<$>)) +import Control.Applicative (Applicative(..), (<$>)) #endif -import Prelude hiding (filter, map, foldl, foldr) +import Control.Applicative (liftA2, liftA3) --- | Map covariantly over a @'WhenMissing' f k x@. -mapWhenMissing :: Functor f => (a -> b) -> WhenMissing f x a -> WhenMissing f x b -mapWhenMissing f q = WhenMissing - { missingSubtree = fmap (map f) . missingSubtree q - , missingKey = \k x -> fmap (forceMaybe . fmap f) $ missingKey q k x} +import Prelude hiding (min, max) --- | Map covariantly over a @'WhenMatched' f k x y@. -mapWhenMatched :: Functor f => (a -> b) -> WhenMatched f x y a -> WhenMatched f x y b -mapWhenMatched f q = WhenMatched - { matchedKey = \k x y -> fmap (forceMaybe . fmap f) $ runWhenMatched q k x y } +import Data.IntMap.Internal +import Data.IntMap.Merge.Internal --- | 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. +(#!), (#) :: (a -> b) -> a -> b +(#!) = ($!) +(#) = ($) + +{-# INLINE fmapMaybe' #-} +fmapMaybe' :: (a -> b) -> Maybe a -> Maybe b +fmapMaybe' _ Nothing = Nothing +fmapMaybe' f (Just x) = Just $! f x + +-- | Map over the entries whose keys are missing from the other map. -- -- @ --- zipWithMaybeMatched :: (k -> x -> y -> Maybe z) --- -> SimpleWhenMatched k x y z +-- mapMissing :: (Key -> a -> b) -> SimpleWhenMissing a b -- @ -zipWithMaybeMatched :: Applicative f - => (Key -> x -> y -> Maybe z) - -> WhenMatched f x y z -zipWithMaybeMatched f = WhenMatched $ - \k x y -> pure $! forceMaybe $! f k x y -{-# INLINE zipWithMaybeMatched #-} - --- | When a key is found in both maps, apply a function to the --- key and values, perform the resulting action, and maybe use --- the result in the merged map. -- --- This is the fundamental 'WhenMatched' tactic. -zipWithMaybeAMatched :: Applicative f - => (Key -> x -> y -> f (Maybe z)) - -> WhenMatched f x y z -zipWithMaybeAMatched f = WhenMatched $ - \ k x y -> forceMaybe <$> f k x y -{-# INLINE zipWithMaybeAMatched #-} +-- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- +-- but @mapMissing@ is somewhat faster. +{-# INLINE mapMissing #-} +mapMissing :: Applicative f => (Key -> a -> b) -> WhenMissing f a b +mapMissing f = mapMissingUKey (\k v -> f (box k) v) --- | When a key is found in both maps, apply a function to the --- key and values to produce an action and use its result in the merged map. -zipWithAMatched :: Applicative f - => (Key -> x -> y -> f z) - -> WhenMatched f x y z -zipWithAMatched f = WhenMatched $ - \ k x y -> (Just $!) <$> f k x y -{-# INLINE zipWithAMatched #-} +-- | Map over the entries whose keys are missing from the other map with a +-- function that takes an unboxed key. Identical in functionality to +-- 'mapMissing'. +mapMissingUKey :: Applicative f => (UKey -> a -> b) -> WhenMissing f a b +mapMissingUKey g = WhenMissing (\k v -> pure (Just $! g k v)) (pure . go g) (pure . go g) (pure . start g) where + start _ Empty = Empty + start f (NonEmpty min minV root) = NonEmpty min #! f (boundUKey min) minV # go f root --- | When a key is found in both maps, apply a function to the --- key and values and use the result in the merged map. --- --- @ --- zipWithMatched :: (k -> x -> y -> z) --- -> SimpleWhenMatched k x y z --- @ -zipWithMatched :: Applicative f - => (Key -> x -> y -> z) -> WhenMatched f x y z -zipWithMatched f = WhenMatched $ - \k x y -> pure $! Just $! f k x y -{-# INLINE zipWithMatched #-} + go :: (UKey -> a -> b) -> Node t a -> Node t b + go _ Tip = Tip + go f (Bin k v l r) = Bin k #! f (boundUKey k) v # go f l # go f r -- | Map over the entries whose keys are missing from the other map, -- optionally removing some. This is the most powerful 'SimpleWhenMissing' -- tactic, but others are usually more efficient. -- -- @ --- mapMaybeMissing :: (k -> x -> Maybe y) -> SimpleWhenMissing k x y +-- mapMaybeMissing :: (Key -> a -> Maybe b) -> SimpleWhenMissing a b -- @ -- -- prop> mapMaybeMissing f = traverseMaybeMissing (\k x -> pure (f k x)) -- -- but @mapMaybeMissing@ uses fewer unnecessary 'Applicative' operations. -mapMaybeMissing :: Applicative f => (Key -> x -> Maybe y) -> WhenMissing f x y -mapMaybeMissing f = WhenMissing - { missingSubtree = \m -> pure $! mapMaybeWithKey f m - , missingKey = \k x -> pure $! forceMaybe $! f k x } {-# INLINE mapMaybeMissing #-} +mapMaybeMissing :: Applicative f => (Key -> a -> Maybe b) -> WhenMissing f a b +mapMaybeMissing f = mapMaybeMissingUKey (\k a -> f (box k) a) --- | Map over the entries whose keys are missing from the other map. +-- | Map over the entries whose keys are missing from the other map using a +-- function taking an unboxed key, optionally removing some. Identical in +-- functionality to 'mapMaybeMissing'. +mapMaybeMissingUKey :: Applicative f => (UKey -> a -> Maybe b) -> WhenMissing f a b +mapMaybeMissingUKey f = WhenMissing (\k v -> case f k v of + Nothing -> pure Nothing + Just !b -> pure (Just b)) (pure . goLKeep) (pure . goRKeep) (pure . start) + where + start Empty = Empty + start (NonEmpty min minV root) = case f (boundUKey min) minV of + Just !minV' -> NonEmpty min minV' (goLKeep root) + Nothing -> goL root + + goLKeep Tip = Tip + goLKeep (Bin max maxV l r) = case f (boundUKey max) maxV of + Just !maxV' -> Bin max maxV' (goLKeep l) (goRKeep r) + Nothing -> binNodeMapL (goLKeep l) (goR r) + + goRKeep Tip = Tip + goRKeep (Bin min minV l r) = case f (boundUKey min) minV of + Just !minV' -> Bin min minV' (goLKeep l) (goRKeep r) + Nothing -> binMapNodeR (goL l) (goRKeep r) + + goL Tip = Empty + goL (Bin max maxV l r) = case f (boundUKey max) maxV of + Just !maxV' -> binL (goL l) (NonEmpty max maxV' (goRKeep r)) + Nothing -> binL (goL l) (goR r) + + goR Tip = Empty + goR (Bin min minV l r) = case f (boundUKey min) minV of + Just !minV' -> binR (NonEmpty min minV' (goLKeep l)) (goR r) + Nothing -> binR (goL l) (goR r) + +-- | 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. -- -- @ --- mapMissing :: (k -> x -> y) -> SimpleWhenMissing k x y +-- zipWithMaybeMatched :: (Key -> a -> b -> Maybe c) +-- -> SimpleWhenMatched a b c -- @ +{-# INLINE zipWithMaybeMatched #-} +zipWithMaybeMatched :: Applicative f => (Key -> a -> b -> Maybe c) -> WhenMatched f a b c +zipWithMaybeMatched f = WhenMatched (\k a b -> case f (box k) a b of + Nothing -> pure Nothing + Just !c -> pure (Just c)) + +-- | When a key is found in both maps, apply a function to the +-- key and values and use the result in the merged map. -- --- prop> mapMissing f = mapMaybeMissing (\k x -> Just $ f k x) +-- @ +-- zipWithMatched :: (Key -> a -> b -> c) +-- -> SimpleWhenMatched a b c +-- @ +{-# INLINE zipWithMatched #-} +zipWithMatched :: Applicative f => (Key -> a -> b -> c) -> WhenMatched f a b c +zipWithMatched f = zipWithMaybeMatched (\k a b -> Just $! f k a b) + +-- | 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. -- --- but @mapMissing@ is somewhat faster. -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 #-} +-- This is the fundamental 'WhenMatched' tactic. +-- +-- @since 0.5.9 +{-# INLINE zipWithMaybeAMatched #-} +zipWithMaybeAMatched + :: (Key -> a -> b -> f (Maybe c)) + -> WhenMatched f a b c +zipWithMaybeAMatched f = WhenMatched (\k a b -> f (box k) a b) + +-- | 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 +{-# INLINE zipWithAMatched #-} +zipWithAMatched + :: Applicative f + => (Key -> a -> b -> f c) + -> WhenMatched f a b c +zipWithAMatched f = zipWithMaybeAMatched (\k a b -> Just <$> f k a b) --- | 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 +-- | Traverse over the entries whose keys are missing from the other +-- map, optionally producing values to put in the result. This is +-- the most powerful 'WhenMissing' tactic, but others are usually -- more efficient. -traverseMaybeMissing :: Applicative f - => (Key -> x -> f (Maybe y)) -> WhenMissing f x y -traverseMaybeMissing f = WhenMissing - { missingSubtree = traverseMaybeWithKey f - , missingKey = \k x -> forceMaybe <$> f k x } +-- +-- @since 0.5.9 {-# INLINE traverseMaybeMissing #-} +traverseMaybeMissing + :: Applicative f => (Key -> a -> f (Maybe b)) -> WhenMissing f a b +traverseMaybeMissing f = WhenMissing + { missingAllL = start + , missingLeft = goL + , missingRight = goR + , missingSingle = f' } + where + f' k a = f (box k) a --- | Traverse over the entries whose keys are missing from the other map. -traverseMissing :: Applicative f - => (Key -> x -> f y) -> WhenMissing f x y -traverseMissing f = WhenMissing - { missingSubtree = traverseWithKey f - , missingKey = \k x -> (Just $!) <$> f k x } + start Empty = pure Empty + start (NonEmpty min minV root) = liftA2 (maybe nodeToMapL (NonEmpty min $!)) (f' (boundUKey min) minV) (goL root) + + goL Tip = pure Tip + goL (Bin max maxV l r) = liftA3 (\l' r' maxV' -> maybe extractBinL (Bin max $!) maxV' l' r') (goL l) (goR r) (f' (boundUKey max) maxV) + + goR Tip = pure Tip + goR (Bin min minV l r) = liftA3 (\minV' l' r' -> maybe extractBinR (Bin min $!) minV' l' r') (f' (boundUKey min) minV) (goL l) (goR r) + +-- | Traverse over the entries whose keys are missing from the other +-- map. +-- +-- @since 0.5.9 {-# INLINE traverseMissing #-} +traverseMissing + :: Applicative f => (Key -> a -> f b) -> WhenMissing f a b +traverseMissing f = WhenMissing + { missingAllL = start + , missingLeft = goL + , missingRight = goR + , missingSingle = \k v -> Just <$> f' k v } + where + f' k a = f (box k) a -forceMaybe :: Maybe a -> Maybe a -forceMaybe Nothing = Nothing -forceMaybe m@(Just !_) = m -{-# INLINE forceMaybe #-} + start Empty = pure Empty + start (NonEmpty min minV root) = liftA2 (NonEmpty min $!) (f' (boundUKey min) minV) (goL root) + + goL Tip = pure Tip + goL (Bin max maxV l r) = liftA3 (\l' r' !maxV' -> Bin max maxV' l' r') (goL l) (goR r) (f' (boundUKey max) maxV) + + goR Tip = pure Tip + goR (Bin min minV l r) = liftA3 (\ !minV' l' r' -> Bin min minV' l' r') (f' (boundUKey min) minV) (goL l) (goR r) + +-- | Map covariantly over a @'WhenMissing' f x@. +-- +-- @since 0.5.9 +{-# INLINE mapWhenMissing #-} +mapWhenMissing :: Functor f => (a -> b) -> WhenMissing f x a -> WhenMissing f x b +mapWhenMissing f miss = WhenMissing + (\k x -> fmap (fmapMaybe' f) (missingSingle miss k x)) + (\l -> fmap (mapNodeStrict f) (missingLeft miss l)) + (\r -> fmap (mapNodeStrict f) (missingRight miss r)) + (\m -> fmap (mapStrict_ f) (missingAllL miss m)) + +-- | Map covariantly over a @'WhenMatched' f x y@. +-- +-- @since 0.5.9 +{-# INLINE mapWhenMatched #-} +mapWhenMatched :: Functor f => (a -> b) -> WhenMatched f x y a -> WhenMatched f x y b +mapWhenMatched f match = WhenMatched (\k x y -> fmap (fmapMaybe' f) (matchedSingle match k x y)) diff --git a/containers/src/Data/IntMap/Strict.hs b/containers/src/Data/IntMap/Strict.hs index ef732c2e2..779abdcea 100644 --- a/containers/src/Data/IntMap/Strict.hs +++ b/containers/src/Data/IntMap/Strict.hs @@ -1,16 +1,21 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP, BangPatterns #-} + +#include "containers.h" + #if !defined(TESTING) && defined(__GLASGOW_HASKELL__) +#if USE_REWRITE_RULES {-# LANGUAGE Trustworthy #-} +#else +{-# LANGUAGE Safe #-} +#endif #endif - -#include "containers.h" ----------------------------------------------------------------------------- -- | -- Module : Data.IntMap.Strict --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 +-- Copyright : Documentation & Interface (c) Daan Leijen 2002 +-- Documentation (c) Andriy Palamarchuk 2008 +-- Documentation & Implementation (c) Jonathan "gereeter" S. 2020 -- License : BSD-style -- Maintainer : libraries@haskell.org -- Portability : portable @@ -67,31 +72,21 @@ -- -- == 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. +-- Operation comments contain the operation time complexity in +-- the Big-O notation . +-- Many operations have a worst-case complexity of /O(min(n,W))/. +-- This means that the operation can become linear in the number of +-- elements with a maximum of /W/ -- the number of bits in an 'Int' +-- (32 or 64). -- +-- Be aware that the 'Functor', 'Traversable' and 'Data' instances +-- are the same as for the "Data.IntMap.Lazy" module, so if they are used +-- on strict maps, the resulting maps will be lazy. ----------------------------------------------------------------------------- --- See the notes at the beginning of Data.IntMap.Internal. - module Data.IntMap.Strict ( -- * Map type -#if !defined(TESTING) - IntMap, Key -- instance Eq,Show -#else - IntMap(..), Key -- instance Eq,Show -#endif + IntMap, Key -- * Construction , empty @@ -143,7 +138,6 @@ module Data.IntMap.Strict ( , size -- * Combine - -- ** Union , union , unionWith @@ -165,7 +159,7 @@ module Data.IntMap.Strict ( -- ** Disjoint , disjoint - -- ** Universal combining function + -- ** Deprecated, unsafe general combining function , mergeWithKey -- * Traversal @@ -202,7 +196,7 @@ module Data.IntMap.Strict ( -- ** Lists , toList --- ** Ordered lists + -- ** Ordered Lists , toAscList , toDescList @@ -213,19 +207,19 @@ module Data.IntMap.Strict ( , withoutKeys , partition , partitionWithKey - , mapMaybe , mapMaybeWithKey , mapEither , mapEitherWithKey - , split , splitLookup , splitRoot -- * Submap - , isSubmapOf, isSubmapOfBy - , isProperSubmapOf, isProperSubmapOfBy + , isSubmapOf + , isSubmapOfBy + , isProperSubmapOf + , isProperSubmapOfBy -- * Min\/Max , lookupMin @@ -245,12 +239,1218 @@ module Data.IntMap.Strict ( , minViewWithKey , maxViewWithKey -#ifdef __GLASGOW_HASKELL__ +#if defined(__GLASGOW_HASKELL__) -- * Debugging , showTree , showTreeWith #endif - ) where +) where + +import Data.IntMap.Internal +#if USE_REWRITE_RULES +import qualified Data.IntMap.Lazy as L +#endif +import qualified Data.IntMap.Merge.Internal as Merge (runWhenMissingAll) +import qualified Data.IntMap.Merge.Strict as Merge (merge, mapMissing, mapMaybeMissing, traverseMissing, zipWithMaybeMatched) +#if defined(__GLASGOW_HASKELL__) +import Data.IntMap.Internal.DeprecatedDebug +#endif + +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (runIdentity) +#else +import Data.IntMap.Merge.Internal (runIdentity) +import Control.Applicative (Applicative(..), (<$>)) +#endif + +import Utils.Containers.Internal.StrictPair (StrictPair(..), toPair) + +import qualified Data.Foldable (Foldable, foldl') +import qualified Data.List (foldl', map) +import qualified Data.IntSet (IntSet, toList) + +import Prelude hiding (foldr, foldl, lookup, null, map, filter, min, max) + +(#!), (#) :: (a -> b) -> a -> b +(#!) = ($!) +(#) = ($) + +wheval :: a -> () +wheval !_ = () + +-- | /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 !v = IntMap (NonEmpty (Bound k) v Tip) + +-- | /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 = insertWithEval wheval const + +-- | /O(min(n,W))/. Insert with a combining function. +-- @'insertWith' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert @f new_value old_value@. +-- +-- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] +-- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" +insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWith = insertWithEval wheval + +-- | /O(min(n,W))/. Insert with a combining function. +-- @'insertWithKey' f key value mp@ +-- will insert the pair (key, value) into @mp@ if key does +-- not exist in the map. If the key does exist, the function will +-- insert @f key new_value old_value@. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] +-- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] +-- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" +insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a +insertWithKey f k = insertWith (f k) k + +-- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@) +-- is a pair where the first element is equal to (@'lookup' k map@) +-- and the second element equal to (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) +-- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) +-- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) +-- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) +insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) +insertLookupWithKey combine !k v = toPair . start + where + start (IntMap Empty) = Nothing :*: IntMap (NonEmpty (Bound k) #! v # Tip) + start (IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> let mv :*: root' = goL (xor k min) min root + in mv :*: IntMap (NonEmpty min minV root') + OutOfBound -> Nothing :*: IntMap (NonEmpty (Bound k) #! v # insertMinL (xor k min) min minV root) + Matched -> Just minV :*: IntMap (NonEmpty (Bound k) #! combine k v minV # root) + + goL !_ _ Tip = Nothing :*: (Bin (Bound k) #! v # Tip # Tip) + goL !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> let mv :*: l' = goL xorCache min l + in mv :*: Bin max maxV l' r + | otherwise -> let mv :*: r' = goR xorCacheMax max r + in mv :*: Bin max maxV l r' + OutOfBound | xor (boundKey max) min < xorCacheMax -> Nothing :*: (Bin (Bound k) #! v # Bin max maxV l r # Tip) + | otherwise -> Nothing :*: (Bin (Bound k) #! v # l # insertMaxR xorCacheMax max maxV r) + Matched -> Just maxV :*: (Bin max #! combine k v maxV # l # r) + where xorCacheMax = xor k max + + goR !_ _ Tip = Nothing :*: (Bin (Bound k) #! v # Tip # Tip) + goR !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> let mv :*: r' = goR xorCache max r + in mv :*: Bin min minV l r' + | otherwise -> let mv :*: l' = goL xorCacheMin min l + in mv :*: Bin min minV l' r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Nothing :*: (Bin (Bound k) #! v # Tip # Bin min minV l r) + | otherwise -> Nothing :*: (Bin (Bound k) #! v # insertMinL xorCacheMin min minV l # r) + Matched -> Just minV :*: (Bin min #! combine k v minV # l # r) + where xorCacheMin = xor k min + +-- | /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 = start + where + start (IntMap Empty) = IntMap Empty + start m@(IntMap (NonEmpty min minV node)) = case compareMinBound k min of + InBound -> IntMap (NonEmpty min minV (goL (xor k min) node)) + OutOfBound -> m + Matched -> IntMap (NonEmpty min #! f minV # node) + + goL !_ Tip = Tip + goL !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goL xorCache l) r + | otherwise -> Bin max maxV l (goR xorCacheMax r) + OutOfBound -> n + Matched -> Bin max #! f maxV # l # r + where xorCacheMax = xor k max + + goR !_ Tip = Tip + goR !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goR xorCache r) + | otherwise -> Bin min minV (goL xorCacheMin l) r + OutOfBound -> n + Matched -> Bin min #! f minV # l # r + where xorCacheMin = xor k min + +-- | /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 = adjust (f k) k + +-- | /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 !k = start + where + start (IntMap Empty) = IntMap Empty + start m@(IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> IntMap (NonEmpty min minV (goL (xor k min) root)) + OutOfBound -> m + Matched -> case f minV of + Nothing -> IntMap (nodeToMapL root) + Just !minV' -> IntMap (NonEmpty min minV' root) + + goL !_ Tip = Tip + goL !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goL xorCache l) r + | otherwise -> Bin max maxV l (goR xorCacheMax r) + OutOfBound -> n + Matched -> case f maxV of + Nothing -> extractBinL l r + Just !maxV' -> Bin max maxV' l r + where xorCacheMax = xor k max + + goR !_ Tip = Tip + goR !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goR xorCache r) + | otherwise -> Bin min minV (goL xorCacheMin l) r + OutOfBound -> n + Matched -> case f minV of + Nothing -> extractBinR l r + Just !minV' -> Bin min minV' l r + where xorCacheMin = xor k min + +-- | /O(min(n,W))/. The expression (@'updateWithKey' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" +updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a +updateWithKey f k = update (f k) k + +-- | /O(min(n,W))/. Lookup and update. +-- The function returns original value, if it is updated. +-- This is different behavior than 'Data.Map.updateLookupWithKey'. +-- Returns the original key value if the map entry is deleted. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")]) +-- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) +-- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") +updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a, IntMap a) +updateLookupWithKey = start + where + start _ !_ (IntMap Empty) = (Nothing, IntMap Empty) + start f !k m@(IntMap (NonEmpty min minV root)) = case compareMinBound k min of + InBound -> let mv :*: root' = goL f k (xor k min) root + in (mv, IntMap (NonEmpty min minV root')) + OutOfBound -> (Nothing, m) + Matched -> (Just minV, case f k minV of + Nothing -> IntMap (nodeToMapL root) + Just !minV' -> IntMap (NonEmpty min minV' root)) + + goL _ !_ !_ Tip = Nothing :*: Tip + goL f !k !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> let mv :*: l' = goL f k xorCache l + in mv :*: Bin max maxV l' r + | otherwise -> let mv :*: r' = goR f k xorCacheMax r + in mv :*: Bin max maxV l r' + OutOfBound -> Nothing :*: n + Matched -> Just maxV :*: case f k maxV of + Nothing -> extractBinL l r + Just !maxV' -> Bin max maxV' l r + where xorCacheMax = xor k max + + goR _ !_ !_ Tip = Nothing :*: Tip + goR f !k !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> let mv :*: r' = goR f k xorCache r + in mv :*: Bin min minV l r' + | otherwise -> let mv :*: l' = goL f k xorCacheMin l + in mv :*: Bin min minV l' r + OutOfBound -> Nothing :*: n + Matched -> Just minV :*: case f k minV of + Nothing -> extractBinR l r + Just !minV' -> Bin min minV' l r + where xorCacheMin = xor k min + +-- | /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 m = case lookup k m of + Nothing -> case f Nothing of + Nothing -> m + Just v -> insert k v m + Just v -> case f (Just v) of + Nothing -> delete k m + Just v' -> insert k v' m + +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, +-- or update a value in an 'IntMap'. In short : @'lookup' k '<$>' 'alterF' f k m = f +-- ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing = do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) = do +-- putStrLn $ "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserResponse2 :: IO (Maybe String) +-- @ +-- +-- 'alterF' is the most general operation for working with an individual +-- key that may or may not be in a given map. +-- +-- Note: 'alterF' is a flipped version of the 'at' combinator from +-- 'Control.Lens.At'. +-- +-- @since 0.5.8 +alterF :: Functor f => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) +alterF f k m = case lookup k m of + Nothing -> fmap (\ret -> case ret of + Nothing -> m + Just v -> insert k v m) (f Nothing) + Just v -> fmap (\ret -> case ret of + Nothing -> delete k m + Just v' -> insert k v' m) (f (Just v)) + +-- | /O(n+m)/. The union with a combining function. +-- +-- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] +unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWith f = unionWithKey (const f) + +-- | /O(n+m)/. The union with a combining function. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] +{-# INLINE unionWithKey #-} +unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithKey combine = unionWithUKey (\k a b -> combine (box k) a b) + +-- | /O(n+m)/. The union with a combining function taking an unboxed key. +-- Identical in functionality to 'unionWithKey'. +unionWithUKey :: (UKey -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a +unionWithUKey = start + where + start _ (IntMap Empty) m2 = m2 + start _ m1 (IntMap Empty) = m1 + start combine (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (NonEmpty min1 minV1 (goL2 combine minV2 min1 root1 min2 root2)) + | min1 > min2 = IntMap (NonEmpty min2 minV2 (goL1 combine minV1 min1 root1 min2 root2)) + | otherwise = IntMap (NonEmpty min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 root1 root2) -- we choose min1 arbitrarily, as min1 == min2 + + goL1 _ minV1 !min1 Tip !_ Tip = Bin (minToMax min1) minV1 Tip Tip + goL1 _ minV1 !min1 !n1 !min2 Tip = insertMinL (xor (boundKey min1) min2) min1 minV1 n1 + goL1 _ minV1 !min1 !n1 !min2 n2@(Bin max2 _ _ _) | boundsDisjoint min1 max2 = unionDisjointL minV1 min2 n2 min1 n1 + goL1 combine minV1 !min1 Tip !min2 !n2 = goInsertL1 combine (boundKey min1) minV1 (xor (boundKey min1) min2) min2 n2 + goL1 combine minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> Bin max2 maxV2 (goL1 combine minV1 min1 n1 min2 l2) r2 -- we choose min1 arbitrarily - we just need something from tree 1 + | max1 > max2 -> Bin max1 maxV1 l2 (goR2 combine maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> Bin max2 maxV2 l2 (goR1 combine maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2) + | otherwise -> Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # l2 # goRFused combine max1 (Bin min1 minV1 l1 r1) r2 -- we choose max1 arbitrarily, as max1 == max2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # goL1 combine minV1 min1 l1 min2 l2 # goRFused combine max1 r1 r2 -- we choose max1 arbitrarily, as max1 == max2 + GT -> Bin max1 maxV1 (goL1 combine minV1 min1 l1 min2 n2) r1 + + goL2 _ minV2 !_ Tip !min2 Tip = Bin (minToMax min2) minV2 Tip Tip + goL2 _ minV2 !min1 Tip !min2 !n2 = insertMinL (xor (boundKey min2) min1) min2 minV2 n2 + goL2 _ minV2 !min1 n1@(Bin max1 _ _ _) !min2 !n2 | boundsDisjoint min2 max1 = unionDisjointL minV2 min1 n1 min2 n2 + goL2 combine minV2 !min1 !n1 !min2 Tip = goInsertL2 combine (boundKey min2) minV2 (xor (boundKey min2) min1) min1 n1 + goL2 combine minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 n2) r1 -- we choose min2 arbitrarily - we just need something from tree 2 + | max1 > max2 -> Bin max1 maxV1 l1 (goR2 combine maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | max1 < max2 -> Bin max2 maxV2 l1 (goR1 combine maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | otherwise -> Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # l1 # goRFused combine max1 r1 (Bin min2 minV2 l2 r2) -- we choose max1 arbitrarily, as max1 == max2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # goL2 combine minV2 min1 l1 min2 l2 # goRFused combine max1 r1 r2 -- we choose max1 arbitrarily, as max1 == max2 + LT -> Bin max2 maxV2 (goL2 combine minV2 min1 n1 min2 l2) r2 + + -- 'goLFused' is called instead of 'goL' if the minimums of the two trees are the same + -- Note that because of this property, the trees cannot be disjoint, so we can skip most of the checks in 'goL' + goLFused _ !_ Tip n2 = n2 + goLFused _ !_ n1 Tip = n1 + goLFused combine min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> Bin max2 maxV2 (goLFused combine min n1 l2) r2 + EQ | max1 > max2 -> Bin max1 maxV1 (goLFused combine min l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> Bin max2 maxV2 (goLFused combine min l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # goLFused combine min l1 l2 # goRFused combine max1 r1 r2 -- we choose max1 arbitrarily, as max1 == max2 + GT -> Bin max1 maxV1 (goLFused combine min l1 n2) r1 + + goR1 _ maxV1 !max1 Tip !_ Tip = Bin (maxToMin max1) maxV1 Tip Tip + goR1 _ maxV1 !max1 !n1 !max2 Tip = insertMaxR (xor (boundKey max1) max2) max1 maxV1 n1 + goR1 _ maxV1 !max1 !n1 !max2 n2@(Bin min2 _ _ _) | boundsDisjoint min2 max1 = unionDisjointR maxV1 max1 n1 max2 n2 + goR1 combine maxV1 !max1 Tip !max2 !n2 = goInsertR1 combine (boundKey max1) maxV1 (xor (boundKey max1) max2) max2 n2 + goR1 combine maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> Bin min2 minV2 l2 (goR1 combine maxV1 max1 n1 max2 r2) -- we choose max1 arbitrarily - we just need something from tree 1 + | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2) r2 + | otherwise -> Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 (Bin max1 maxV1 l1 r1) l2 # r2 -- we choose min1 arbitrarily, as min1 == min2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2 # goR1 combine maxV1 max1 r1 max2 r2 -- we choose min1 arbitrarily, as min1 == min2 + GT -> Bin min1 minV1 l1 (goR1 combine maxV1 max1 r1 max2 n2) + + goR2 _ maxV2 !_ Tip !max2 Tip = Bin (maxToMin max2) maxV2 Tip Tip + goR2 _ maxV2 !max1 Tip !max2 !n2 = insertMaxR (xor (boundKey max2) max1) max2 maxV2 n2 + goR2 _ maxV2 !max1 n1@(Bin min1 _ _ _) !max2 !n2 | boundsDisjoint min1 max2 = unionDisjointR maxV2 max2 n2 max1 n1 + goR2 combine maxV2 !max1 !n1 !max2 Tip = goInsertR2 combine (boundKey max2) maxV2 (xor (boundKey max2) max1) max1 n1 + goR2 combine maxV2 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> Bin min1 minV1 l1 (goR2 combine maxV2 max1 r1 max2 n2) -- we choose max2 arbitrarily - we just need something from tree 2 + | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | otherwise -> Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 (Bin max2 maxV2 l2 r2) # r1 -- we choose min1 arbitrarily, as min1 == min2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | otherwise -> Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2 # goR2 combine maxV2 max1 r1 max2 r2 -- we choose min1 arbitrarily, as min1 == min2 + LT -> Bin min2 minV2 l2 (goR2 combine maxV2 max1 n1 max2 r2) + + -- 'goRFused' is called instead of 'goR' if the minimums of the two trees are the same + -- Note that because of this property, the trees cannot be disjoint, so we can skip most of the checks in 'goR' + goRFused _ !_ Tip n2 = n2 + goRFused _ !_ n1 Tip = n1 + goRFused combine max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> Bin min2 minV2 l2 (goRFused combine max n1 r2) + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goRFused combine max r1 r2) + | min1 > min2 -> Bin min2 minV2 (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max r1 r2) + | otherwise -> Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2 # goRFused combine max r1 r2 -- we choose min1 arbitrarily, as min1 == min2 + GT -> Bin min1 minV1 l1 (goRFused combine max r1 n2) + + goInsertL1 _ k v !_ _ Tip = Bin (Bound k) #! v # Tip # Tip + goInsertL1 combine k v !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goInsertL1 combine k v xorCache min l) r + | otherwise -> Bin max maxV l (goInsertR1 combine k v xorCacheMax max r) + OutOfBound | xor (boundKey max) min < xorCacheMax -> Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Bin max #! combine (unbox k) v maxV # l # r + where xorCacheMax = xor k max + + goInsertR1 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertR1 combine k v !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goInsertR1 combine k v xorCache max r) + | otherwise -> Bin min minV (goInsertL1 combine k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Bin min #! combine (unbox k) v minV # l # r + where xorCacheMin = xor k min + + goInsertL2 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertL2 combine k v !xorCache min (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goInsertL2 combine k v xorCache min l) r + | otherwise -> Bin max maxV l (goInsertR2 combine k v xorCacheMax max r) + OutOfBound | xor (boundKey max) min < xorCacheMax -> Bin (Bound k) v (Bin max maxV l r) Tip + | otherwise -> Bin (Bound k) v l (insertMaxR xorCacheMax max maxV r) + Matched -> Bin max #! combine (unbox k) maxV v # l # r + where xorCacheMax = xor k max + + goInsertR2 _ k v !_ _ Tip = Bin (Bound k) v Tip Tip + goInsertR2 combine k v !xorCache max (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goInsertR2 combine k v xorCache max r) + | otherwise -> Bin min minV (goInsertL2 combine k v xorCacheMin min l) r + OutOfBound | xor (boundKey min) max < xorCacheMin -> Bin (Bound k) v Tip (Bin min minV l r) + | otherwise -> Bin (Bound k) v (insertMinL xorCacheMin min minV l) r + Matched -> Bin min #! combine (unbox k) minV v # l # r + where xorCacheMin = xor k min + +-- | 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 :: Data.Foldable.Foldable f => (a -> a -> a) -> f (IntMap a) -> IntMap a +unionsWith f = Data.Foldable.foldl' (unionWith f) empty + +-- | /O(n+m)/. Difference with a combining function. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) +-- > == singleton 3 "b:B" +differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWith f = differenceWithKey (const f) + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). +-- If it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) +-- > == singleton 3 "3:b|B" +{-# INLINE differenceWithKey #-} +differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWithKey combine = differenceWithUKey (\k a b -> combine (box k) a b) + +-- | /O(n+m)/. The difference with a combining function taking an unboxed +-- key. Identical in functionality to 'differenceWithKey'. +differenceWithUKey :: (UKey -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a +differenceWithUKey = start + where + start _ (IntMap Empty) !_ = IntMap Empty + start _ !m (IntMap Empty) = m + start combine (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (NonEmpty min1 minV1 (goL2 combine minV2 min1 root1 min2 root2)) + | min1 > min2 = IntMap (goL1 combine minV1 min1 root1 min2 root2) + | otherwise = case combine (boundUKey min1) minV1 minV2 of + Nothing -> IntMap (goLFused combine min1 root1 root2) + Just !minV1' -> IntMap (NonEmpty min1 minV1' (goLFusedKeep combine min1 root1 root2)) + + goL1 combine minV1 !min1 Tip !min2 !n2 = goLookupL combine (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL1 _ minV1 !min1 !n1 !_ Tip = NonEmpty min1 minV1 n1 + goL1 _ minV1 !min1 n1@(Bin _ _ _ _) !_ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = NonEmpty min1 minV1 n1 + goL1 combine minV1 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL1 combine minV1 min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> r2lMap $ NonEmpty max1 maxV1 (goR2 combine maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2) + | max1 < max2 -> r2lMap $ goR1 combine maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> r2lMap $ goRFused combine max1 (Bin min1 minV1 l1 r1) r2 + Just !maxV1' -> r2lMap $ NonEmpty max1 maxV1' (goRFusedKeep combine max1 (Bin min1 minV1 l1 r1) r2) + EQ | max1 > max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (NonEmpty max1 maxV1 (goR2 combine maxV2 max1 r1 max2 r2)) + | max1 < max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> binL (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max1 r1 r2) + Just !maxV1' -> binL (goL1 combine minV1 min1 l1 min2 l2) (NonEmpty max1 maxV1' (goRFusedKeep combine max1 r1 r2)) + GT -> binL (goL1 combine minV1 min1 l1 min2 n2) (NonEmpty max1 maxV1 r1) + + goL2 _ _ !_ Tip !_ !_ = Tip + goL2 combine minV2 !min1 !n1 !min2 Tip = goDeleteL combine (boundKey min2) minV2 (xor (boundKey min2) min1) n1 + goL2 _ _ !_ n1@(Bin max1 _ _ _) !min2 (Bin _ _ _ _) | boundsDisjoint min2 max1 = n1 + goL2 combine minV2 !min1 n1@(Bin max1 maxV1 l1 r1) !min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goL2 combine minV2 min1 n1 min2 l2 + EQ | max1 > max2 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> goL2 combine minV2 min1 l1 min2 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goL2 combine minV2 min1 l1 min2 l2) r' + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> case goRFused combine max1 r1 r2 of + Empty -> goL2 combine minV2 min1 l1 min2 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goL2 combine minV2 min1 l1 min2 l2) r' + Just !maxV1' -> Bin max1 maxV1' (goL2 combine minV2 min1 l1 min2 l2) (goRFusedKeep combine max1 r1 r2) + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> Bin max1 maxV1 (goL2 combine minV2 min1 l1 min2 n2) r1 -- min2 is arbitrary here - we just need something from tree 2 + | max1 > max2 -> Bin max1 maxV1 l1 (goR2 combine maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2)) + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2) of + Empty -> l1 + NonEmpty max' maxV' r' -> Bin max' maxV' l1 r' + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> case goRFused combine max1 r1 (Bin min2 minV2 l2 r2) of + Empty -> l1 + NonEmpty max' maxV' r' -> Bin max' maxV' l1 r' + Just !maxV1' -> Bin max1 maxV1' l1 (goRFusedKeep combine max1 r1 (Bin min2 minV2 l2 r2)) + + goLFused _ !_ Tip !_ = Empty + goLFused _ !_ (Bin max1 maxV1 l1 r1) Tip = case deleteMinL max1 maxV1 l1 r1 of + NE min' minV' n' -> NonEmpty min' minV' n' + goLFused combine !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFused combine min n1 l2 + EQ | max1 > max2 -> binL (goLFused combine min l1 l2) (NonEmpty max1 maxV1 (goR2 combine maxV2 max1 r1 max2 r2)) + | max1 < max2 -> binL (goLFused combine min l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> binL (goLFused combine min l1 l2) (goRFused combine max1 r1 r2) -- we choose max1 arbitrarily, as max1 == max2 + Just !maxV1' -> binL (goLFused combine min l1 l2) (NonEmpty max1 maxV1' (goRFusedKeep combine max1 r1 r2)) + GT -> binL (goLFused combine min l1 n2) (NonEmpty max1 maxV1 r1) + + goLFusedKeep _ !_ !n1 Tip = n1 + goLFusedKeep _ !_ Tip !_ = Tip + goLFusedKeep combine !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFusedKeep combine min n1 l2 + EQ | max1 > max2 -> Bin max1 maxV1 (goLFusedKeep combine min l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> goLFusedKeep combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFusedKeep combine min l1 l2) r' + | otherwise -> case combine (boundUKey max1) maxV1 maxV2 of + Nothing -> case goRFused combine max1 r1 r2 of -- we choose max1 arbitrarily, as max1 == max2 + Empty -> goLFusedKeep combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFusedKeep combine min l1 l2) r' + Just !maxV1' -> Bin max1 maxV1' (goLFusedKeep combine min l1 l2) (goRFusedKeep combine max1 r1 r2) + GT -> Bin max1 maxV1 (goLFusedKeep combine min l1 n2) r1 + + goR1 combine maxV1 !max1 Tip !max2 !n2 = goLookupR combine (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR1 _ maxV1 !max1 !n1 !_ Tip = NonEmpty max1 maxV1 n1 + goR1 _ maxV1 !max1 n1@(Bin _ _ _ _) !_ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = NonEmpty max1 maxV1 n1 + goR1 combine maxV1 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR1 combine maxV1 max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> l2rMap $ NonEmpty min1 minV1 (goL2 combine minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2) + | min1 > min2 -> l2rMap $ goL1 combine minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> l2rMap $ goLFused combine min1 (Bin max1 maxV1 l1 r1) l2 + Just !minV1' -> l2rMap $ NonEmpty min1 minV1' (goLFusedKeep combine min1 (Bin max1 maxV1 l1 r1) l2) + EQ | min1 < min2 -> binR (NonEmpty min1 minV1 (goL2 combine minV2 min1 l1 min2 l2)) (goR1 combine maxV1 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> binR (goLFused combine min1 l1 l2) (goR1 combine maxV1 max1 r1 max2 r2) + Just !minV1' -> binR (NonEmpty min1 minV1' (goLFusedKeep combine min1 l1 l2)) (goR1 combine maxV1 max1 r1 max2 r2) + GT -> binR (NonEmpty min1 minV1 l1) (goR1 combine maxV1 max1 r1 max2 n2) + + goR2 _ _ !_ Tip !_ !_ = Tip + goR2 combine maxV2 !max1 !n1 !max2 Tip = goDeleteR combine (boundKey max2) maxV2 (xor (boundKey max2) max1) n1 + goR2 _ _ !_ n1@(Bin min1 _ _ _) !max2 (Bin _ _ _ _) | boundsDisjoint min1 max2 = n1 + goR2 combine maxV2 !max1 n1@(Bin min1 minV1 l1 r1) !max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goR2 combine maxV2 max1 n1 max2 r2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> goR2 combine maxV2 max1 r1 max2 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goR2 combine maxV2 max1 r1 max2 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> case goLFused combine min1 l1 l2 of + Empty -> goR2 combine maxV2 max1 r1 max2 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goR2 combine maxV2 max1 r1 max2 r2) + Just !minV1' -> Bin min1 minV1' (goLFusedKeep combine min1 l1 l2) (goR2 combine maxV2 max1 r1 max2 r2) + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> Bin min1 minV1 l1 (goR2 combine maxV2 max1 r1 max2 n2) -- max2 is arbitrary here - we just need something from tree 2 + | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2)) r1 + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2) of + Empty -> r1 + NonEmpty min' minV' l' -> Bin min' minV' l' r1 + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> case goLFused combine min1 l1 (Bin max2 maxV2 l2 r2) of + Empty -> r1 + NonEmpty min' minV' l' -> Bin min' minV' l' r1 + Just !minV1' -> Bin min1 minV1' (goLFusedKeep combine min1 l1 (Bin max2 maxV2 l2 r2)) r1 + + goRFused _ !_ Tip !_ = Empty + goRFused _ !_ (Bin min1 minV1 l1 r1) Tip = case deleteMaxR min1 minV1 l1 r1 of + NE max' maxV' n' -> NonEmpty max' maxV' n' + goRFused combine !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFused combine max n1 r2 + EQ | min1 < min2 -> binR (NonEmpty min1 minV1 (goL2 combine minV2 min1 l1 min2 l2)) (goRFused combine max r1 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goRFused combine max r1 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of + Nothing -> binR (goLFused combine min1 l1 l2) (goRFused combine max r1 r2) -- we choose min1 arbitrarily, as min1 == min2 + Just !minV1' -> binR (NonEmpty min1 minV1' (goLFusedKeep combine min1 l1 l2)) (goRFused combine max r1 r2) + GT -> binR (NonEmpty min1 minV1 l1) (goRFused combine max r1 n2) + + goRFusedKeep _ !_ !n1 Tip = n1 + goRFusedKeep _ !_ Tip !_ = Tip + goRFusedKeep combine !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFusedKeep combine max n1 r2 + EQ | min1 < min2 -> Bin min1 minV1 (goL2 combine minV2 min1 l1 min2 l2) (goRFusedKeep combine max r1 r2) + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> goRFusedKeep combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFusedKeep combine max r1 r2) + | otherwise -> case combine (boundUKey min1) minV1 minV2 of -- we choose min1 arbitrarily, as min1 == min2 + Nothing -> case goLFused combine min1 l1 l2 of + Empty -> goRFusedKeep combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFusedKeep combine max r1 r2) + Just !minV1' -> Bin min1 minV1' (goLFusedKeep combine min1 l1 l2) (goRFusedKeep combine max r1 r2) + GT -> Bin min1 minV1 l1 (goRFusedKeep combine max r1 n2) + + goLookupL _ !k v !_ Tip = NonEmpty (Bound k) v Tip + goLookupL combine !k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL combine k v xorCache l + | otherwise -> goLookupR combine k v xorCacheMax r + OutOfBound -> NonEmpty (Bound k) v Tip + Matched -> case combine (unbox k) v maxV of + Nothing -> Empty + Just !v' -> NonEmpty (Bound k) v' Tip + where xorCacheMax = xor k max + + goLookupR _ !k v !_ Tip = NonEmpty (Bound k) v Tip + goLookupR combine !k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR combine k v xorCache r + | otherwise -> goLookupL combine k v xorCacheMin l + OutOfBound -> NonEmpty (Bound k) v Tip + Matched -> case combine (unbox k) v minV of + Nothing -> Empty + Just !v' -> NonEmpty (Bound k) v' Tip + where xorCacheMin = xor k min -import Data.IntMap.Strict.Internal -import Prelude () + goDeleteL _ !_ _ !_ Tip = Tip + goDeleteL combine !k v !xorCache n@(Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> Bin max maxV (goDeleteL combine k v xorCache l) r + | otherwise -> Bin max maxV l (goDeleteR combine k v xorCacheMax r) + OutOfBound -> n + Matched -> case combine (unbox k) maxV v of + Nothing -> extractBinL l r + Just !maxV' -> Bin max maxV' l r + where xorCacheMax = xor k max + + goDeleteR _ !_ _ !_ Tip = Tip + goDeleteR combine !k v !xorCache n@(Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> Bin min minV l (goDeleteR combine k v xorCache r) + | otherwise -> Bin min minV (goDeleteL combine k v xorCacheMin l) r + OutOfBound -> n + Matched -> case combine (unbox k) minV v of + Nothing -> extractBinR l r + Just !minV' -> Bin min minV' l r + where xorCacheMin = xor k min + +-- | /O(n+m)/. The intersection with a combining function. +-- +-- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" +intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWith f = intersectionWithKey (const f) + +-- | /O(n+m)/. The intersection with a combining function. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" +{-# INLINE intersectionWithKey #-} +intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWithKey combine = intersectionWithUKey (\k a b -> combine (box k) a b) + +-- | /O(n+m)/. The intersection with a combining function taking an unboxed +-- key. Identical in functionality to 'intersectionWithKey'. +intersectionWithUKey :: (UKey -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c +intersectionWithUKey = start + where + start _ (IntMap Empty) !_ = IntMap Empty + start _ !_ (IntMap Empty) = IntMap Empty + start combine (IntMap (NonEmpty min1 minV1 root1)) (IntMap (NonEmpty min2 minV2 root2)) + | min1 < min2 = IntMap (goL2 combine minV2 min1 root1 min2 root2) + | min1 > min2 = IntMap (goL1 combine minV1 min1 root1 min2 root2) + | otherwise = IntMap (NonEmpty min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 root1 root2) -- we choose min1 arbitrarily, as min1 == min2 + + -- TODO: This scheme might produce lots of unnecessary l2r and r2l calls. This should be rectified. + + goL1 _ _ !_ !_ !_ Tip = Empty + goL1 combine minV1 min1 Tip min2 n2 = goLookupL1 combine (boundKey min1) minV1 (xor (boundKey min1) min2) n2 + goL1 _ _ min1 (Bin _ _ _ _) _ (Bin max2 _ _ _) | boundsDisjoint min1 max2 = Empty + goL1 combine minV1 min1 n1@(Bin max1 maxV1 l1 r1) min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey min1) min2 < xor (boundKey min1) max2 -> goL1 combine minV1 min1 n1 min2 l2 -- min1 is arbitrary here - we just need something from tree 1 + | max1 > max2 -> r2lMap $ goR2 combine maxV2 max1 (Bin min1 minV1 l1 r1) max2 r2 + | max1 < max2 -> r2lMap $ goR1 combine maxV1 max1 (Bin min1 minV1 l1 r1) max2 r2 + | otherwise -> r2lMap $ NonEmpty max1 #! combine (boundUKey max1) maxV1 maxV2 # goRFused combine max1 (Bin min1 minV1 l1 r1) r2 + EQ | max1 > max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> binL (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> r2lMap (NonEmpty max1 #! combine (boundUKey max1) maxV1 maxV2 # goRFused combine max1 r1 r2) + NonEmpty min' minV' l' -> NonEmpty min' minV' (Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # l' # goRFused combine max1 r1 r2) + GT -> goL1 combine minV1 min1 l1 min2 n2 + + goL2 _ _ !_ Tip !_ !_ = Empty + goL2 combine minV2 min1 n1 min2 Tip = goLookupL2 combine (boundKey min2) minV2 (xor (boundKey min2) min1) n1 + goL2 _ _ _ (Bin max1 _ _ _) min2 (Bin _ _ _ _) | boundsDisjoint min2 max1 = Empty + goL2 combine minV2 min1 n1@(Bin max1 maxV1 l1 r1) min2 n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goL2 combine minV2 min1 n1 min2 l2 + EQ | max1 > max2 -> binL (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | max1 < max2 -> binL (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case goL2 combine minV2 min1 l1 min2 l2 of + Empty -> r2lMap (NonEmpty max1 #! combine (boundUKey max1) maxV1 maxV2 # goRFused combine max1 r1 r2) + NonEmpty min' minV' l' -> NonEmpty min' minV' (Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # l' # goRFused combine max1 r1 r2) + GT | xor (boundKey min2) min1 < xor (boundKey min2) max1 -> goL2 combine minV2 min1 l1 min2 n2 -- min2 is arbitrary here - we just need something from tree 2 + | max1 > max2 -> r2lMap $ goR2 combine maxV2 max1 r1 max2 (Bin min2 minV2 l2 r2) + | max1 < max2 -> r2lMap $ goR1 combine maxV1 max1 r1 max2 (Bin min2 minV2 l2 r2) + | otherwise -> r2lMap $ NonEmpty max1 #! combine (boundUKey max1) maxV1 maxV2 # goRFused combine max1 r1 (Bin min2 minV2 l2 r2) + + goLFused _ !_ Tip !_ = Tip + goLFused _ !_ !_ Tip = Tip + goLFused combine !min n1@(Bin max1 maxV1 l1 r1) n2@(Bin max2 maxV2 l2 r2) = case compareMSB (xorBounds min max1) (xorBounds min max2) of + LT -> goLFused combine min n1 l2 + EQ | max1 > max2 -> case goR2 combine maxV2 max1 r1 max2 r2 of + Empty -> goLFused combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFused combine min l1 l2) r' + | max1 < max2 -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> goLFused combine min l1 l2 + NonEmpty max' maxV' r' -> Bin max' maxV' (goLFused combine min l1 l2) r' + | otherwise -> Bin max1 #! combine (boundUKey max1) maxV1 maxV2 # goLFused combine min l1 l2 # goRFused combine max1 r1 r2 -- we choose max1 arbitrarily, as max1 == max2 + GT -> goLFused combine min l1 n2 + + goR1 _ _ !_ !_ !_ Tip = Empty + goR1 combine maxV1 max1 Tip max2 n2 = goLookupR1 combine (boundKey max1) maxV1 (xor (boundKey max1) max2) n2 + goR1 _ _ max1 (Bin _ _ _ _) _ (Bin min2 _ _ _) | boundsDisjoint min2 max1 = Empty + goR1 combine maxV1 max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT | xor (boundKey max1) min2 > xor (boundKey max1) max2 -> goR1 combine maxV1 max1 n1 max2 r2 -- max1 is arbitrary here - we just need something from tree 1 + | min1 < min2 -> l2rMap $ goL2 combine minV2 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | min1 > min2 -> l2rMap $ goL1 combine minV1 min1 (Bin max1 maxV1 l1 r1) min2 l2 + | otherwise -> l2rMap $ NonEmpty min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 (Bin max1 maxV1 l1 r1) l2 + EQ | min1 < min2 -> binR (goL2 combine minV2 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goR1 combine maxV1 max1 r1 max2 r2) + | otherwise -> case goR1 combine maxV1 max1 r1 max2 r2 of + Empty -> l2rMap (NonEmpty min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2) + NonEmpty max' maxV' r' -> NonEmpty max' maxV' (Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2 # r') + GT -> goR1 combine maxV1 max1 r1 max2 n2 + + goR2 _ _ !_ Tip !_ !_ = Empty + goR2 combine maxV2 max1 n1 max2 Tip = goLookupR2 combine (boundKey max2) maxV2 (xor (boundKey max2) max1) n1 + goR2 _ _ _ (Bin min1 _ _ _) max2 (Bin _ _ _ _) | boundsDisjoint min1 max2 = Empty + goR2 combine maxV2 max1 n1@(Bin min1 minV1 l1 r1) max2 n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max1) (xorBounds min2 max2) of + LT -> goR2 combine maxV2 max1 n1 max2 r2 + EQ | min1 < min2 -> binR (goL2 combine minV2 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | min1 > min2 -> binR (goL1 combine minV1 min1 l1 min2 l2) (goR2 combine maxV2 max1 r1 max2 r2) + | otherwise -> case goR2 combine maxV2 max1 r1 max2 r2 of + Empty -> l2rMap (NonEmpty min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2) + NonEmpty max' maxV' r' -> NonEmpty max' maxV' (Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2 # r') + GT | xor (boundKey max2) min1 > xor (boundKey max2) max1 -> goR2 combine maxV2 max1 r1 max2 n2 -- max2 is arbitrary here - we just need something from tree 2 + | min1 < min2 -> l2rMap $ goL2 combine minV2 min1 l1 min2 (Bin max2 maxV2 l2 r2) + | min1 > min2 -> l2rMap $ goL1 combine minV1 min1 l1 min2 (Bin max2 maxV2 l2 r2) + | otherwise -> l2rMap $ NonEmpty min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 (Bin max2 maxV2 l2 r2) + + goRFused _ !_ Tip !_ = Tip + goRFused _ !_ !_ Tip = Tip + goRFused combine !max n1@(Bin min1 minV1 l1 r1) n2@(Bin min2 minV2 l2 r2) = case compareMSB (xorBounds min1 max) (xorBounds min2 max) of + LT -> goRFused combine max n1 r2 + EQ | min1 < min2 -> case goL2 combine minV2 min1 l1 min2 l2 of + Empty -> goRFused combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFused combine max r1 r2) + | min1 > min2 -> case goL1 combine minV1 min1 l1 min2 l2 of + Empty -> goRFused combine max r1 r2 + NonEmpty min' minV' l' -> Bin min' minV' l' (goRFused combine max r1 r2) + | otherwise -> Bin min1 #! combine (boundUKey min1) minV1 minV2 # goLFused combine min1 l1 l2 # goRFused combine max r1 r2 -- we choose max1 arbitrarily, as max1 == max2 + GT -> goRFused combine max r1 n2 + + goLookupL1 _ !_ _ !_ Tip = Empty + goLookupL1 combine k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL1 combine k v xorCache l + | otherwise -> goLookupR1 combine k v xorCacheMax r + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) #! combine (unbox k) v maxV # Tip + where xorCacheMax = xor k max + + goLookupR1 _ !_ _ !_ Tip = Empty + goLookupR1 combine k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR1 combine k v xorCache r + | otherwise -> goLookupL1 combine k v xorCacheMin l + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) #! combine (unbox k) v minV # Tip + where xorCacheMin = xor k min + + goLookupL2 _ !_ _ !_ Tip = Empty + goLookupL2 combine k v !xorCache (Bin max maxV l r) = case compareMaxBound k max of + InBound | xorCache < xorCacheMax -> goLookupL2 combine k v xorCache l + | otherwise -> goLookupR2 combine k v xorCacheMax r + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) #! combine (unbox k) maxV v # Tip + where xorCacheMax = xor k max + + goLookupR2 _ !_ _ !_ Tip = Empty + goLookupR2 combine k v !xorCache (Bin min minV l r) = case compareMinBound k min of + InBound | xorCache < xorCacheMin -> goLookupR2 combine k v xorCache r + | otherwise -> goLookupL2 combine k v xorCacheMin l + OutOfBound -> Empty + Matched -> NonEmpty (Bound k) #! combine (unbox k) minV v # Tip + where xorCacheMin = xor k min + +-- | /O(n+m)/. An unsafe general combining function. +-- +-- WARNING: This function can produce corrupt maps and its results +-- may depend on the internal structures of its inputs. Users should +-- prefer 'merge' or 'mergeA'. This function is also significantly slower +-- than 'merge'. +-- +-- When 'mergeWithKey' is given three arguments, it is inlined to the call +-- site. You should therefore use 'mergeWithKey' only to define custom +-- combining functions. For example, you could define 'unionWithKey', +-- 'differenceWithKey' and 'intersectionWithKey' as +-- +-- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 +-- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 +-- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 +-- +-- When calling @'mergeWithKey' combine only1 only2@, a function combining two +-- '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@, +-- @'filterWithKey' f@, or @'mapMaybeWithKey' f@ could be used for any @f@. +mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) -> IntMap a -> IntMap b -> IntMap c +mergeWithKey matched miss1 miss2 = Merge.merge (Merge.mapMaybeMissing (single miss1)) (Merge.mapMaybeMissing (single miss2)) (Merge.zipWithMaybeMatched matched) where + single miss k v = case miss (IntMap (NonEmpty (Bound k) v Tip)) of + IntMap Empty -> Nothing + IntMap (NonEmpty _ v' _) -> Just v' + +-- | /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 (IntMap m) = IntMap (mapStrict_ f m) + +-- | /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 = runIdentity . Merge.runWhenMissingAll (Merge.mapMissing f) + +#if USE_REWRITE_RULES +-- Pay close attention to strictness here. We need to force the +-- intermediate result for @'map' f . 'map' g@, and we need to refrain +-- from forcing it for @'map' f . 'L.map' g@, etc. +-- +-- TODO Consider writing RULES for things like @'L.map' f ('map' g m)@. We'd +-- need a new function for this, and we'd have to pay attention to simplifier +-- phases. Somthing like +-- +-- > lsmap :: (b -> c) -> (a -> b) -> Node t a -> Node t c +-- > lsmap _ _ Tip = Tip +-- > lsmap f g (Bin bound value l r) = let !gvalue = g value +-- > in Bin bound (f gvalue) (lsmap f g l) (lsmap f g r) +{-# NOINLINE[1] map #-} +{-# NOINLINE[1] mapWithKey #-} +{-# RULES +"map/map" forall f g m . map f (map g m) = map (\v -> f $! g v) m +"map/mapL" forall f g m . map f (mapLazy g m) = map (f . g) m +"mapWithKey/map" forall f g m . mapWithKey f (map g m) = mapWithKey (\k v -> f k $! g v) m +"mapWithKey/mapL" forall f g m . mapWithKey f (mapLazy g m) = mapWithKey (\k -> f k . g) m +"map/mapWithKey" forall f g m . map f (mapWithKey g m) = mapWithKey (\k v -> f $! g k v) m +"map/mapWithKeyL" forall f g m . map f (L.mapWithKey g m) = mapWithKey (\k -> f . g k) m +"mapWithKey/mapWithKey" forall f g m . mapWithKey f (mapWithKey g m) = mapWithKey (\k v -> f k $! g k v) m +"mapWithKey/mapWithKeyL" forall f g m . mapWithKey f (L.mapWithKey g m) = mapWithKey (\k -> f k . g k) m + #-} +#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 +{-# INLINE traverseWithKey #-} +traverseWithKey :: Applicative f => (Key -> a -> f b) -> IntMap a -> f (IntMap b) +traverseWithKey f = Merge.runWhenMissingAll (Merge.traverseMissing f) + +-- | /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 = start + where + start a (IntMap Empty) = (a, IntMap Empty) + start a (IntMap (NonEmpty min minV root)) = + let (a', !minV') = f a (boundKey min) minV + (a'', root') = goL root a' + in (a'', IntMap (NonEmpty min minV' root')) + + goL Tip a = (a, Tip) + goL (Bin max maxV l r) a = + let (a', l') = goL l a + (a'', r') = goR r a' + (a''', !maxV') = f a'' (boundKey max) maxV + in (a''', Bin max maxV' l' r') + + goR Tip a = (a, Tip) + goR (Bin min minV l r) a = + let (a', !minV') = f a (boundKey min) minV + (a'', l') = goL l a' + (a''', r') = goR r a'' + in (a''', Bin min minV' l' r') + +-- | /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 = start + where + start a (IntMap Empty) = (a, IntMap Empty) + start a (IntMap (NonEmpty min minV root)) = + let (a', root') = goL root a + (a'', !minV') = f a' (boundKey min) minV + in (a'', IntMap (NonEmpty min minV' root')) + + goL Tip a = (a, Tip) + goL (Bin max maxV l r) a = + let (a', !maxV') = f a (boundKey max) maxV + (a'', r') = goR r a' + (a''', l') = goL l a'' + in (a''', Bin max maxV' l' r') + + goR Tip a = (a, Tip) + goR (Bin min minV l r) a = + let (a', r') = goR r a + (a'', l') = goL l a' + (a''', !minV') = f a'' (boundKey min) minV + in (a''', Bin min minV' l' r') + +-- | /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 = foldlWithKey' (\m k a -> insert (f k) a m) empty + +-- | /O(n*min(n,W))/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" +mapKeysWith :: (a -> a -> a) -> (Key -> Key) -> IntMap a -> IntMap a +mapKeysWith combine f = foldlWithKey' (\m k a -> insertWith combine (f k) a m) empty + +-- | /O(n*min(n,W))/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has slightly better performance than 'mapKeys'. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList [(5,"a"), (3,"b")]) == fromList [(6, "b"), (10, "a")] +mapKeysMonotonic :: (Key -> Key) -> IntMap a -> IntMap a +mapKeysMonotonic = mapKeys + +-- | /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) -> Data.IntSet.IntSet -> IntMap a +fromSet f = fromDistinctAscList . Data.List.map (\k -> (k, f k)) . Data.IntSet.toList + +-- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs. +fromList :: [(Key, a)] -> IntMap a +fromList = Data.List.foldl' (\t (k, a) -> insert k a t) empty + +-- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "ab"), (5, "cba")] +-- > fromListWith (++) [] == empty +fromListWith :: (a -> a -> a) -> [(Key, a)] -> IntMap a +fromListWith f = Data.List.foldl' (\t (k, a) -> insertWith f k a t) empty + +-- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromListWithKey f [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"c")] == fromList [(3, "3:a|b"), (5, "5:c|5:b|a")] +-- > fromListWithKey f [] == empty +fromListWithKey :: (Key -> a -> a -> a) -> [(Key, a)] -> IntMap a +fromListWithKey f = Data.List.foldl' (\t (k, a) -> insertWithKey f k a t) empty + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +fromAscList :: [(Key, a)] -> IntMap a +fromAscList = start where + start [] = IntMap Empty + start ((min, minV) : rest) = IntMap (go min minV rest StackBase) + + -- The strictness here is tricky, mimicking old behavior. Only the last value + -- of a given key (the one being inserted into the map) is evaluated strictly. + go !k v [] !stk = completeBuildStack (Bound k) v Tip stk + go !k v ((next, nextV) : rest) !stk + | next == k = nextV `seq` go k nextV rest stk + | otherwise = v `seq` go next nextV rest (pushBuildStack (xor k (Bound next)) k v Tip stk) + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order. +-- +-- > fromAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +-- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")] +fromAscListWith :: (a -> a -> a) -> [(Key, a)] -> IntMap a +fromAscListWith f = fromAscListWithKey (const f) + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order, with a combining function on equal keys. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > fromAscListWithKey f [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "5:b|a")] +fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key, a)] -> IntMap a +fromAscListWithKey f = start where + start [] = IntMap Empty + start ((min, minV) : rest) = IntMap (go min minV rest StackBase) + + go !k v [] !stk = completeBuildStack (Bound k) v Tip stk + go !k v ((next, nextV) : rest) !stk + | next == k = let !v' = f k nextV v in go k v' rest stk + | otherwise = v `seq` go next nextV rest (pushBuildStack (xor k (Bound next)) k v Tip stk) + +-- | /O(n)/. Build a map from a list of key\/value pairs where +-- the keys are in ascending order and all distinct. +-- /The precondition (input list is strictly ascending) is not checked./ +-- +-- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] +fromDistinctAscList :: [(Key, a)] -> IntMap a +fromDistinctAscList = start where + start [] = IntMap Empty + start ((min, !minV) : rest) = IntMap (go min minV rest StackBase) + + go !k !v [] !stk = completeBuildStack (Bound k) v Tip stk + go !k !v ((next, !nextV) : rest) !stk = go next nextV rest (pushBuildStack (xor k (Bound next)) k v Tip stk) + +-- | /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 (const f) + +-- | /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" +{-# INLINE mapMaybeWithKey #-} +mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b +mapMaybeWithKey f = runIdentity . Merge.runWhenMissingAll (Merge.mapMaybeMissing f) + +-- | /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 = mapEitherWithKey (const f) + +-- | /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")]) +{-# INLINE mapEitherWithKey #-} +mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) +mapEitherWithKey f = mapEitherWithUKey (\k a -> f (box k) a) + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results with +-- a mapping function taking unboxed keys. Identical in functionality to +-- 'mapEitherWithKey'. +mapEitherWithUKey :: (UKey -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) +mapEitherWithUKey func = start + where + start (IntMap Empty) = (IntMap Empty, IntMap Empty) + start (IntMap (NonEmpty min minV root)) = case func (boundUKey min) minV of + Left !v -> let t :*: f = goTrueL root + in (IntMap (NonEmpty min v t), IntMap f) + Right !v -> let t :*: f = goFalseL root + in (IntMap t, IntMap (NonEmpty min v f)) + + goTrueL Tip = Tip :*: Empty + goTrueL (Bin max maxV l r) = case func (boundUKey max) maxV of + Left !v -> let tl :*: fl = goTrueL l + tr :*: fr = goTrueR r + in Bin max v tl tr :*: binL fl fr + Right !v -> let tl :*: fl = goTrueL l + tr :*: fr = goFalseR r + in binNodeMapL tl tr :*: binL fl (NonEmpty max v fr) + + goTrueR Tip = Tip :*: Empty + goTrueR (Bin min minV l r) = case func (boundUKey min) minV of + Left !v -> let tl :*: fl = goTrueL l + tr :*: fr = goTrueR r + in Bin min v tl tr :*: binR fl fr + Right !v -> let tl :*: fl = goFalseL l + tr :*: fr = goTrueR r + in binMapNodeR tl tr :*: binR (NonEmpty min v fl) fr + + goFalseL Tip = Empty :*: Tip + goFalseL (Bin max maxV l r) = case func (boundUKey max) maxV of + Left !v -> let tl :*: fl = goFalseL l + tr :*: fr = goTrueR r + in binL tl (NonEmpty max v tr) :*: binNodeMapL fl fr + Right !v -> let tl :*: fl = goFalseL l + tr :*: fr = goFalseR r + in binL tl tr :*: Bin max v fl fr + + goFalseR Tip = Empty :*: Tip + goFalseR (Bin min minV l r) = case func (boundUKey min) minV of + Left !v -> let tl :*: fl = goTrueL l + tr :*: fr = goFalseR r + in binR (NonEmpty min v tl) tr :*: binMapNodeR fl fr + Right !v -> let tl :*: fl = goFalseL l + tr :*: fr = goFalseR r + in binR tl tr :*: Bin min v fl fr + +-- | /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))/. 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. +-- +-- > 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 _ (IntMap Empty) = IntMap Empty +updateMinWithKey f (IntMap (NonEmpty min minV root)) = IntMap $ case f (boundKey min) minV of + Nothing -> nodeToMapL root + Just !minV' -> NonEmpty min minV' root + +-- | /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 _ (IntMap Empty) = IntMap Empty +updateMaxWithKey f (IntMap (NonEmpty min minV Tip)) = IntMap $ case f (boundKey min) minV of + Nothing -> Empty + Just !minV' -> NonEmpty min minV' Tip +updateMaxWithKey f (IntMap (NonEmpty min minV (Bin max maxV l r))) = IntMap . NonEmpty min minV $ case f (boundKey max) maxV of + Nothing -> extractBinL l r + Just !maxV' -> Bin max maxV' l r diff --git a/containers/src/Data/IntMap/Strict/Internal.hs b/containers/src/Data/IntMap/Strict/Internal.hs deleted file mode 100644 index 2b695c0fe..000000000 --- a/containers/src/Data/IntMap/Strict/Internal.hs +++ /dev/null @@ -1,1209 +0,0 @@ -{-# LANGUAGE CPP #-} -{-# LANGUAGE BangPatterns #-} -{-# LANGUAGE PatternGuards #-} - -#include "containers.h" - ------------------------------------------------------------------------------ --- | --- Module : Data.IntMap.Strict.Internal --- Copyright : (c) Daan Leijen 2002 --- (c) Andriy Palamarchuk 2008 --- License : BSD-style --- Maintainer : libraries@haskell.org --- Portability : portable --- --- --- = Finite Int Maps (strict interface) --- --- The @'IntMap' v@ type represents a finite map (sometimes called a dictionary) --- from key of type @Int@ to values of type @v@. --- --- Each function in this module is careful to force values before installing --- them in an 'IntMap'. This is usually more efficient when laziness is not --- necessary. When laziness /is/ required, use the functions in --- "Data.IntMap.Lazy". --- --- In particular, the functions in this module obey the following law: --- --- - If all values stored in all maps in the arguments are in WHNF, then all --- values stored in all maps in the results will be in WHNF once those maps --- are evaluated. --- --- For a walkthrough of the most commonly used functions see the --- . --- --- This module is intended to be imported qualified, to avoid name clashes with --- Prelude functions: --- --- > import Data.IntMap.Strict (IntMap) --- > import qualified Data.IntMap.Strict as IntMap --- --- 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. --- --- --- == Detailed performance information --- --- The amortized running time is given for each operation, with /n/ referring to --- the number of entries in the map and /W/ referring to the number of bits in --- an 'Int' (32 or 64). --- --- Benchmarks comparing "Data.IntMap.Strict" with other dictionary --- implementations can be found at https://github.com/haskell-perf/dictionaries. --- --- --- == Warning --- --- The 'IntMap' type is shared between the lazy and strict modules, meaning that --- the same 'IntMap' value can be passed to functions in both modules. This --- means that the 'Functor', 'Traversable' and 'Data.Data.Data' instances are --- the same as for the "Data.IntMap.Lazy" module, so if they are used the --- resulting map may contain suspended values (thunks). --- --- --- == 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. --- ------------------------------------------------------------------------------ - --- See the notes at the beginning of Data.IntMap.Internal. - -module Data.IntMap.Strict.Internal ( - -- * Map type -#if !defined(TESTING) - IntMap, Key -- instance Eq,Show -#else - IntMap(..), Key -- instance Eq,Show -#endif - - -- * Construction - , empty - , singleton - , fromSet - - -- ** From Unordered Lists - , fromList - , fromListWith - , fromListWithKey - - -- ** From Ascending Lists - , fromAscList - , fromAscListWith - , fromAscListWithKey - , fromDistinctAscList - - -- * Insertion - , insert - , insertWith - , insertWithKey - , insertLookupWithKey - - -- * Deletion\/Update - , delete - , adjust - , adjustWithKey - , update - , updateWithKey - , updateLookupWithKey - , alter - , alterF - - -- * Query - -- ** Lookup - , lookup - , (!?) - , (!) - , findWithDefault - , member - , notMember - , lookupLT - , lookupGT - , lookupLE - , lookupGE - - -- ** Size - , null - , size - - -- * Combine - - -- ** Union - , union - , unionWith - , unionWithKey - , unions - , unionsWith - - -- ** Difference - , difference - , (\\) - , differenceWith - , differenceWithKey - - -- ** Intersection - , intersection - , intersectionWith - , intersectionWithKey - - -- ** Disjoint - , disjoint - - -- ** Universal combining function - , mergeWithKey - - -- * Traversal - -- ** Map - , map - , mapWithKey - , traverseWithKey - , traverseMaybeWithKey - , mapAccum - , mapAccumWithKey - , mapAccumRWithKey - , mapKeys - , mapKeysWith - , mapKeysMonotonic - - -- * Folds - , foldr - , foldl - , foldrWithKey - , foldlWithKey - , foldMapWithKey - - -- ** Strict folds - , foldr' - , foldl' - , foldrWithKey' - , foldlWithKey' - - -- * Conversion - , elems - , keys - , assocs - , keysSet - - -- ** Lists - , toList - --- ** Ordered lists - , toAscList - , toDescList - - -- * Filter - , filter - , filterWithKey - , restrictKeys - , withoutKeys - , partition - , partitionWithKey - - , 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 - -#ifdef __GLASGOW_HASKELL__ - -- * Debugging - , showTree - , showTreeWith -#endif - ) where - -import Prelude hiding (lookup,map,filter,foldr,foldl,null) - -import Data.Bits -import qualified Data.IntMap.Internal as L -import Data.IntMap.Internal - ( IntMap (..) - , Key - , mask - , branchMask - , nomatch - , zero - , natFromInt - , intFromNat - , bin - , binCheckLeft - , binCheckRight - , link - , linkWithMask - - , (\\) - , (!) - , (!?) - , empty - , assocs - , filter - , filterWithKey - , findMin - , findMax - , foldMapWithKey - , foldr - , foldl - , foldr' - , foldl' - , foldlWithKey - , foldrWithKey - , foldlWithKey' - , foldrWithKey' - , keysSet - , mergeWithKey' - , delete - , deleteMin - , deleteMax - , deleteFindMax - , deleteFindMin - , difference - , elems - , intersection - , disjoint - , isProperSubmapOf - , isProperSubmapOfBy - , isSubmapOf - , isSubmapOfBy - , lookup - , lookupLE - , lookupGE - , lookupLT - , lookupGT - , lookupMin - , lookupMax - , minView - , maxView - , minViewWithKey - , maxViewWithKey - , keys - , mapKeys - , mapKeysMonotonic - , member - , notMember - , null - , partition - , partitionWithKey - , restrictKeys - , size - , split - , splitLookup - , splitRoot - , toAscList - , toDescList - , toList - , union - , unions - , withoutKeys - ) -#ifdef __GLASGOW_HASKELL__ -import Data.IntMap.Internal.DeprecatedDebug (showTree, showTreeWith) -#endif -import qualified Data.IntSet.Internal as IntSet -import Utils.Containers.Internal.BitUtil -import Utils.Containers.Internal.StrictPair -#if !MIN_VERSION_base(4,8,0) -import Data.Functor((<$>)) -#endif -import Control.Applicative (Applicative (..), liftA2) -import qualified Data.Foldable as Foldable -#if !MIN_VERSION_base(4,8,0) -import Data.Foldable (Foldable()) -#endif - -{-------------------------------------------------------------------- - Query ---------------------------------------------------------------------} - --- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@ --- returns the value at key @k@ or returns @def@ when the key is not an --- element of the map. --- --- > findWithDefault 'x' 1 (fromList [(5,'a'), (3,'b')]) == 'x' --- > findWithDefault 'x' 5 (fromList [(5,'a'), (3,'b')]) == 'a' - --- See IntMap.Internal.Note: Local 'go' functions and capturing] -findWithDefault :: a -> Key -> IntMap a -> a -findWithDefault def !k = go - where - go (Bin p m l r) | nomatch k p m = def - | zero k m = go l - | otherwise = go r - go (Tip kx x) | k == kx = x - | otherwise = def - go Nil = def - -{-------------------------------------------------------------------- - Construction ---------------------------------------------------------------------} --- | /O(1)/. A map of one element. --- --- > singleton 1 'a' == fromList [(1, 'a')] --- > size (singleton 1 'a') == 1 - -singleton :: Key -> a -> IntMap a -singleton k !x - = 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 = - case t of - Bin p m l r - | nomatch k p m -> link k (Tip k x) p t - | zero k m -> Bin p m (insert k x l) r - | otherwise -> Bin p m l (insert k x r) - Tip ky _ - | k==ky -> Tip k x - | otherwise -> link k (Tip k x) ky t - Nil -> Tip k x - --- right-biased insertion, used by 'union' --- | /O(min(n,W))/. Insert with a combining function. --- @'insertWith' f key value mp@ --- will insert the pair (key, value) into @mp@ if key does --- not exist in the map. If the key does exist, the function will --- insert @f new_value old_value@. --- --- > insertWith (++) 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "xxxa")] --- > insertWith (++) 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] --- > insertWith (++) 5 "xxx" empty == singleton 5 "xxx" - -insertWith :: (a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWith f k x t - = insertWithKey (\_ x' y' -> f x' y') k x t - --- | /O(min(n,W))/. Insert with a combining function. --- @'insertWithKey' f key value mp@ --- will insert the pair (key, value) into @mp@ if key does --- not exist in the map. If the key does exist, the function will --- insert @f key new_value old_value@. --- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > insertWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:xxx|a")] --- > insertWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a"), (7, "xxx")] --- > insertWithKey f 5 "xxx" empty == singleton 5 "xxx" --- --- If the key exists in the map, this function is lazy in @value@ but strict --- in the result of @f@. - -insertWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> IntMap a -insertWithKey f !k x t = - case t of - Bin p m l r - | nomatch k p m -> link k (singleton k x) p t - | zero k m -> Bin p m (insertWithKey f k x l) r - | otherwise -> Bin p m l (insertWithKey f k x r) - Tip ky y - | k==ky -> Tip k $! f k x y - | otherwise -> link k (singleton k x) ky t - Nil -> singleton k x - --- | /O(min(n,W))/. The expression (@'insertLookupWithKey' f k x map@) --- is a pair where the first element is equal to (@'lookup' k map@) --- and the second element equal to (@'insertWithKey' f k x map@). --- --- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value --- > insertLookupWithKey f 5 "xxx" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:xxx|a")]) --- > insertLookupWithKey f 7 "xxx" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "xxx")]) --- > insertLookupWithKey f 5 "xxx" empty == (Nothing, singleton 5 "xxx") --- --- This is how to define @insertLookup@ using @insertLookupWithKey@: --- --- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t --- > insertLookup 5 "x" (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "x")]) --- > insertLookup 7 "x" (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a"), (7, "x")]) - -insertLookupWithKey :: (Key -> a -> a -> a) -> Key -> a -> IntMap a -> (Maybe a, IntMap a) -insertLookupWithKey f0 !k0 x0 t0 = toPair $ go f0 k0 x0 t0 - where - go f k x t = - case t of - Bin p m l r - | nomatch k p m -> Nothing :*: link k (singleton k x) p t - | zero k m -> let (found :*: l') = go f k x l in (found :*: Bin p m l' r) - | otherwise -> let (found :*: r') = go f k x r in (found :*: Bin p m l r') - Tip ky y - | k==ky -> (Just y :*: (Tip k $! f k x y)) - | otherwise -> (Nothing :*: link k (singleton k x) ky t) - Nil -> Nothing :*: (singleton k x) - - -{-------------------------------------------------------------------- - Deletion ---------------------------------------------------------------------} --- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not --- a member of the map, the original map is returned. --- --- > adjust ("new " ++) 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] --- > adjust ("new " ++) 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > adjust ("new " ++) 7 empty == empty - -adjust :: (a -> a) -> Key -> IntMap a -> IntMap a -adjust f k m - = adjustWithKey (\_ x -> f x) k m - --- | /O(min(n,W))/. Adjust a value at a specific key. When the key is not --- a member of the map, the original map is returned. --- --- > let f key x = (show key) ++ ":new " ++ x --- > adjustWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] --- > adjustWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > adjustWithKey f 7 empty == empty - -adjustWithKey :: (Key -> a -> a) -> Key -> IntMap a -> IntMap a -adjustWithKey f !k t = - case t of - Bin p m l r - | nomatch k p m -> t - | zero k m -> Bin p m (adjustWithKey f k l) r - | otherwise -> Bin p m l (adjustWithKey f k r) - Tip ky y - | k==ky -> Tip ky $! f k y - | otherwise -> t - Nil -> Nil - --- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ --- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is --- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. --- --- > let f x = if x == "a" then Just "new a" else Nothing --- > update f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "new a")] --- > update f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > update f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -update :: (a -> Maybe a) -> Key -> IntMap a -> IntMap a -update f - = updateWithKey (\_ x -> f x) - --- | /O(min(n,W))/. The expression (@'update' f k map@) updates the value @x@ --- at @k@ (if it is in the map). If (@f k x@) is 'Nothing', the element is --- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateWithKey f 5 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "5:new a")] --- > updateWithKey f 7 (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "a")] --- > updateWithKey f 3 (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -updateWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> IntMap a -updateWithKey f !k t = - case t of - Bin p m l r - | nomatch k p m -> t - | zero k m -> binCheckLeft p m (updateWithKey f k l) r - | otherwise -> binCheckRight p m l (updateWithKey f k r) - Tip ky y - | k==ky -> case f k y of - Just !y' -> Tip ky y' - Nothing -> Nil - | otherwise -> t - Nil -> Nil - --- | /O(min(n,W))/. Lookup and update. --- The function returns original value, if it is updated. --- This is different behavior than 'Data.Map.updateLookupWithKey'. --- Returns the original key value if the map entry is deleted. --- --- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing --- > updateLookupWithKey f 5 (fromList [(5,"a"), (3,"b")]) == (Just "a", fromList [(3, "b"), (5, "5:new a")]) --- > updateLookupWithKey f 7 (fromList [(5,"a"), (3,"b")]) == (Nothing, fromList [(3, "b"), (5, "a")]) --- > updateLookupWithKey f 3 (fromList [(5,"a"), (3,"b")]) == (Just "b", singleton 5 "a") - -updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> IntMap a -> (Maybe a,IntMap a) -updateLookupWithKey f0 !k0 t0 = toPair $ go f0 k0 t0 - where - go f k t = - case t of - Bin p m l r - | nomatch k p m -> (Nothing :*: t) - | zero k m -> let (found :*: l') = go f k l in (found :*: binCheckLeft p m l' r) - | otherwise -> let (found :*: r') = go f k r in (found :*: binCheckRight p m l r') - Tip ky y - | k==ky -> case f k y of - Just !y' -> (Just y :*: Tip ky y') - Nothing -> (Just y :*: Nil) - | otherwise -> (Nothing :*: t) - Nil -> (Nothing :*: Nil) - - - --- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. --- 'alter' can be used to insert, delete, or update a value in an 'IntMap'. --- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. -alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a -alter f !k t = - case t of - Bin p m l r - | nomatch k p m -> case f Nothing of - Nothing -> t - Just !x -> link k (Tip k x) p t - | zero k m -> binCheckLeft p m (alter f k l) r - | otherwise -> binCheckRight p m l (alter f k r) - Tip ky y - | k==ky -> case f (Just y) of - Just !x -> Tip ky x - Nothing -> Nil - | otherwise -> case f Nothing of - Just !x -> link k (Tip k x) ky t - Nothing -> t - Nil -> case f Nothing of - Just !x -> Tip k x - Nothing -> Nil - --- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at --- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, --- or update a value in an 'IntMap'. In short : @'lookup' k <$> 'alterF' f k m = f --- ('lookup' k m)@. --- --- Example: --- --- @ --- interactiveAlter :: Int -> IntMap String -> IO (IntMap String) --- interactiveAlter k m = alterF f k m where --- f Nothing = do --- putStrLn $ show k ++ --- " was not found in the map. Would you like to add it?" --- getUserResponse1 :: IO (Maybe String) --- f (Just old) = do --- putStrLn $ "The key is currently bound to " ++ show old ++ --- ". Would you like to change or delete it?" --- getUserResponse2 :: IO (Maybe String) --- @ --- --- 'alterF' is the most general operation for working with an individual --- key that may or may not be in a given map. - --- Note: 'alterF' is a flipped version of the 'at' combinator from --- 'Control.Lens.At'. --- --- @since 0.5.8 - -alterF :: Functor f - => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a) --- This implementation was modified from 'Control.Lens.At'. -alterF f k m = (<$> f mv) $ \fres -> - case fres of - Nothing -> maybe m (const (delete k m)) mv - Just !v' -> insert k v' m - where mv = lookup k m - - -{-------------------------------------------------------------------- - Union ---------------------------------------------------------------------} --- | 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(n+m)/. The union with a combining function. --- --- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")] - -unionWith :: (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWith f m1 m2 - = unionWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. The union with a combining function. --- --- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value --- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")] - -unionWithKey :: (Key -> a -> a -> a) -> IntMap a -> IntMap a -> IntMap a -unionWithKey f m1 m2 - = mergeWithKey' Bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) id id m1 m2 - -{-------------------------------------------------------------------- - Difference ---------------------------------------------------------------------} - --- | /O(n+m)/. Difference with a combining function. --- --- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing --- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")]) --- > == singleton 3 "b:B" - -differenceWith :: (a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a -differenceWith f m1 m2 - = differenceWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. Difference with a combining function. When two equal keys are --- encountered, the combining function is applied to the key and both values. --- If it returns 'Nothing', the element is discarded (proper set difference). --- If it returns (@'Just' y@), the element is updated with a new value @y@. --- --- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing --- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")]) --- > == singleton 3 "3:b|B" - -differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMap a -differenceWithKey f m1 m2 - = mergeWithKey f id (const Nil) m1 m2 - -{-------------------------------------------------------------------- - Intersection ---------------------------------------------------------------------} - --- | /O(n+m)/. The intersection with a combining function. --- --- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA" - -intersectionWith :: (a -> b -> c) -> IntMap a -> IntMap b -> IntMap c -intersectionWith f m1 m2 - = intersectionWithKey (\_ x y -> f x y) m1 m2 - --- | /O(n+m)/. The intersection with a combining function. --- --- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar --- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A" - -intersectionWithKey :: (Key -> a -> b -> c) -> IntMap a -> IntMap b -> IntMap c -intersectionWithKey f m1 m2 - = mergeWithKey' bin (\(Tip k1 x1) (Tip _k2 x2) -> Tip k1 $! f k1 x1 x2) (const Nil) (const Nil) m1 m2 - -{-------------------------------------------------------------------- - MergeWithKey ---------------------------------------------------------------------} - --- | /O(n+m)/. A high-performance universal combining function. Using --- 'mergeWithKey', all combining functions can be defined without any loss of --- efficiency (with exception of 'union', 'difference' and 'intersection', --- where sharing of some nodes is lost with 'mergeWithKey'). --- --- Please make sure you know what is going on when using 'mergeWithKey', --- otherwise you can be surprised by unexpected code growth or even --- corruption of the data structure. --- --- When 'mergeWithKey' is given three arguments, it is inlined to the call --- site. You should therefore use 'mergeWithKey' only to define your custom --- combining functions. For example, you could define 'unionWithKey', --- 'differenceWithKey' and 'intersectionWithKey' as --- --- > myUnionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) id id m1 m2 --- > myDifferenceWithKey f m1 m2 = mergeWithKey f id (const empty) m1 m2 --- > myIntersectionWithKey f m1 m2 = mergeWithKey (\k x1 x2 -> Just (f k x1 x2)) (const empty) (const empty) m1 m2 --- --- When calling @'mergeWithKey' combine only1 only2@, a function combining two --- 'IntMap's is created, such that --- --- * if a key is present in both maps, it is passed with both corresponding --- values to the @combine@ function. Depending on the result, the key is either --- present in the result with specified value, or is left out; --- --- * a nonempty subtree present only in the first map is passed to @only1@ and --- the output is added to the result; --- --- * a nonempty subtree present only in the second map is passed to @only2@ and --- the output is added to the result. --- --- The @only1@ and @only2@ methods /must return a map with a subset (possibly empty) of the keys of the given map/. --- The values can be modified arbitrarily. Most common variants of @only1@ and --- @only2@ are 'id' and @'const' 'empty'@, but for example @'map' f@ or --- @'filterWithKey' f@ could be used for any @f@. - -mergeWithKey :: (Key -> a -> b -> Maybe c) -> (IntMap a -> IntMap c) -> (IntMap b -> IntMap c) - -> IntMap a -> IntMap b -> IntMap c -mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2 - where -- We use the lambda form to avoid non-exhaustive pattern matches warning. - combine = \(Tip k1 x1) (Tip _k2 x2) -> case f k1 x1 x2 of Nothing -> Nil - Just !x -> Tip k1 x - {-# INLINE combine #-} -{-# INLINE mergeWithKey #-} - -{-------------------------------------------------------------------- - Min\/Max ---------------------------------------------------------------------} - --- | /O(log n)/. Update the value at the minimal key. --- --- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"3:b"), (5,"a")] --- > updateMinWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a -updateMinWithKey f t = - case t of Bin p m l r | m < 0 -> binCheckRight p m l (go f r) - _ -> go f t - where - go f' (Bin p m l r) = binCheckLeft p m (go f' l) r - go f' (Tip k y) = case f' k y of - Just !y' -> Tip k y' - Nothing -> Nil - go _ Nil = error "updateMinWithKey Nil" - --- | /O(log n)/. Update the value at the maximal key. --- --- > updateMaxWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3,"b"), (5,"5:a")] --- > updateMaxWithKey (\ _ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" - -updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a -updateMaxWithKey f t = - case t of Bin p m l r | m < 0 -> binCheckLeft p m (go f l) r - _ -> go f t - where - go f' (Bin p m l r) = binCheckRight p m l (go f' r) - go f' (Tip k y) = case f' k y of - Just !y' -> Tip k y' - Nothing -> Nil - go _ Nil = error "updateMaxWithKey Nil" - --- | /O(log n)/. Update the value at the maximal key. --- --- > updateMax (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "b"), (5, "Xa")] --- > updateMax (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 3 "b" - -updateMax :: (a -> Maybe a) -> IntMap a -> IntMap a -updateMax f = updateMaxWithKey (const f) - --- | /O(log n)/. Update the value at the minimal key. --- --- > updateMin (\ a -> Just ("X" ++ a)) (fromList [(5,"a"), (3,"b")]) == fromList [(3, "Xb"), (5, "a")] --- > updateMin (\ _ -> Nothing) (fromList [(5,"a"), (3,"b")]) == singleton 5 "a" - -updateMin :: (a -> Maybe a) -> IntMap a -> IntMap a -updateMin f = updateMinWithKey (const f) - - -{-------------------------------------------------------------------- - Mapping ---------------------------------------------------------------------} --- | /O(n)/. Map a function over all values in the map. --- --- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")] - -map :: (a -> b) -> IntMap a -> IntMap b -map f = go - where - go (Bin p m l r) = Bin p m (go l) (go r) - go (Tip k x) = Tip k $! f x - go Nil = Nil - -#ifdef __GLASGOW_HASKELL__ -{-# NOINLINE [1] map #-} -{-# RULES -"map/map" forall f g xs . map f (map g xs) = map (\x -> f $! g x) xs -"map/mapL" forall f g xs . map f (L.map g xs) = map (\x -> f (g x)) xs - #-} -#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 m l r -> Bin p m (mapWithKey f l) (mapWithKey f r) - Tip k x -> Tip k $! f k x - Nil -> Nil - -#ifdef __GLASGOW_HASKELL__ --- Pay close attention to strictness here. We need to force the --- intermediate result for map f . map g, and we need to refrain --- from forcing it for map f . L.map g, etc. --- --- TODO Consider moving map and mapWithKey to IntMap.Internal so we can write --- non-orphan RULES for things like L.map f (map g xs). We'd need a new function --- for this, and we'd have to pay attention to simplifier phases. Something like --- --- lsmap :: (b -> c) -> (a -> b) -> IntMap a -> IntMap c --- lsmap _ _ Nil = Nil --- lsmap f g (Tip k x) = let !gx = g x in Tip k (f gx) --- lsmap f g (Bin p m l r) = Bin p m (lsmap f g l) (lsmap f g r) -{-# 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/mapWithKeyL" forall f g xs . mapWithKey f (L.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 -"mapWithKey/mapL" forall f g xs . mapWithKey f (L.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 -"map/mapWithKeyL" forall f g xs . map f (L.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) = (\ !v' -> Tip k v') <$> f k v - go (Bin p m l r) - | m < 0 = liftA2 (flip (Bin p m)) (go r) (go l) - | otherwise = liftA2 (Bin p m) (go l) (go r) -{-# INLINE traverseWithKey #-} - --- | /O(n)/. Traverse keys\/values and collect the 'Just' results. -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 m l r) - | m < 0 = liftA2 (flip (bin p m)) (go r) (go l) - | otherwise = liftA2 (bin p m) (go l) (go r) - --- | /O(n)/. The function @'mapAccum'@ threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a b = (a ++ b, b ++ "X") --- > mapAccum f "Everything: " (fromList [(5,"a"), (3,"b")]) == ("Everything: ba", fromList [(3, "bX"), (5, "aX")]) - -mapAccum :: (a -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccum f = mapAccumWithKey (\a' _ x -> f a' x) - --- | /O(n)/. The function @'mapAccumWithKey'@ threads an accumulating --- argument through the map in ascending order of keys. --- --- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") --- > mapAccumWithKey f "Everything:" (fromList [(5,"a"), (3,"b")]) == ("Everything: 3-b 5-a", fromList [(3, "bX"), (5, "aX")]) - -mapAccumWithKey :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccumWithKey f a t - = mapAccumL f a t - --- | /O(n)/. The function @'mapAccumL'@ threads an accumulating --- argument through the map in ascending order of keys. Strict in --- the accumulating argument and the both elements of the --- result of the function. -mapAccumL :: (a -> Key -> b -> (a,c)) -> a -> IntMap b -> (a,IntMap c) -mapAccumL f0 a0 t0 = toPair $ go f0 a0 t0 - where - go f a t - = case t of - Bin p m l r - | m < 0 -> - let (a1 :*: r') = go f a r - (a2 :*: l') = go f a1 l - in (a2 :*: Bin p m l' r') - | otherwise -> - let (a1 :*: l') = go f a l - (a2 :*: r') = go f a1 r - in (a2 :*: Bin p m l' r') - Tip k x -> let !(a',!x') = f a k x in (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 f0 a0 t0 = toPair $ go f0 a0 t0 - where - go f a t - = case t of - Bin p m l r - | m < 0 -> - let (a1 :*: l') = go f a l - (a2 :*: r') = go f a1 r - in (a2 :*: Bin p m l' r') - | otherwise -> - let (a1 :*: r') = go f a r - (a2 :*: l') = go f a1 l - in (a2 :*: Bin p m l' r') - Tip k x -> let !(a',!x') = f a k x in (a' :*: Tip k x') - Nil -> (a :*: Nil) - --- | /O(n*log n)/. --- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. --- --- The size of the result may be smaller if @f@ maps two or more distinct --- keys to the same new key. In this case the associated values will be --- combined using @c@. --- --- > mapKeysWith (++) (\ _ -> 1) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 1 "cdab" --- > mapKeysWith (++) (\ _ -> 3) (fromList [(1,"b"), (2,"a"), (3,"d"), (4,"c")]) == singleton 3 "cdab" - -mapKeysWith :: (a -> a -> a) -> (Key->Key) -> IntMap a -> IntMap a -mapKeysWith c f = fromListWith c . foldrWithKey (\k x xs -> (f k, x) : xs) [] - -{-------------------------------------------------------------------- - Filter ---------------------------------------------------------------------} --- | /O(n)/. Map values and collect the 'Just' results. --- --- > let f x = if x == "a" then Just "new a" else Nothing --- > mapMaybe f (fromList [(5,"a"), (3,"b")]) == singleton 5 "new a" - -mapMaybe :: (a -> Maybe b) -> IntMap a -> IntMap b -mapMaybe f = mapMaybeWithKey (\_ x -> f x) - --- | /O(n)/. Map keys\/values and collect the 'Just' results. --- --- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing --- > mapMaybeWithKey f (fromList [(5,"a"), (3,"b")]) == singleton 3 "key : 3" - -mapMaybeWithKey :: (Key -> a -> Maybe b) -> IntMap a -> IntMap b -mapMaybeWithKey f (Bin p m l r) - = bin p m (mapMaybeWithKey f l) (mapMaybeWithKey f r) -mapMaybeWithKey f (Tip k x) = case f k x of - Just !y -> Tip k y - Nothing -> Nil -mapMaybeWithKey _ Nil = Nil - --- | /O(n)/. Map values and separate the 'Left' and 'Right' results. --- --- > let f a = if a < "c" then Left a else Right a --- > mapEither f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (fromList [(3,"b"), (5,"a")], fromList [(1,"x"), (7,"z")]) --- > --- > mapEither (\ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (empty, fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) - -mapEither :: (a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) -mapEither f m - = mapEitherWithKey (\_ x -> f x) m - --- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. --- --- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) --- > mapEitherWithKey f (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (fromList [(1,2), (3,6)], fromList [(5,"aa"), (7,"zz")]) --- > --- > mapEitherWithKey (\_ a -> Right a) (fromList [(5,"a"), (3,"b"), (1,"x"), (7,"z")]) --- > == (empty, fromList [(1,"x"), (3,"b"), (5,"a"), (7,"z")]) - -mapEitherWithKey :: (Key -> a -> Either b c) -> IntMap a -> (IntMap b, IntMap c) -mapEitherWithKey f0 t0 = toPair $ go f0 t0 - where - go f (Bin p m l r) - = bin p m l1 r1 :*: bin p m l2 r2 - where - (l1 :*: l2) = go f l - (r1 :*: r2) = go f r - go f (Tip k x) = case f k x of - Left !y -> (Tip k y :*: Nil) - Right !z -> (Nil :*: Tip k z) - go _ Nil = (Nil :*: Nil) - -{-------------------------------------------------------------------- - Conversions ---------------------------------------------------------------------} - --- | /O(n)/. Build a map from a set of keys and a function which for each key --- computes its value. --- --- > fromSet (\k -> replicate k 'a') (Data.IntSet.fromList [3, 5]) == fromList [(5,"aaaaa"), (3,"aaa")] --- > fromSet undefined Data.IntSet.empty == empty - -fromSet :: (Key -> a) -> IntSet.IntSet -> IntMap a -fromSet _ IntSet.Nil = Nil -fromSet f (IntSet.Bin p m l r) = Bin p m (fromSet f l) (fromSet f r) -fromSet f (IntSet.Tip kx bm) = buildTree f kx bm (IntSet.suffixBitMask + 1) - where -- This is slightly complicated, as we to convert the dense - -- representation of IntSet into tree representation of IntMap. - -- - -- We are given a nonzero bit mask 'bmask' of 'bits' bits with prefix 'prefix'. - -- We split bmask into halves corresponding to left and right subtree. - -- If they are both nonempty, we create a Bin node, otherwise exactly - -- one of them is nonempty and we construct the IntMap from that half. - buildTree g !prefix !bmask bits = case bits of - 0 -> Tip prefix $! g prefix - _ -> case intFromNat ((natFromInt bits) `shiftRL` 1) of - bits2 | bmask .&. ((1 `shiftLL` bits2) - 1) == 0 -> - buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2 - | (bmask `shiftRL` bits2) .&. ((1 `shiftLL` bits2) - 1) == 0 -> - buildTree g prefix bmask bits2 - | otherwise -> - Bin prefix bits2 (buildTree g prefix bmask bits2) (buildTree g (prefix + bits2) (bmask `shiftRL` bits2) bits2) - -{-------------------------------------------------------------------- - Lists ---------------------------------------------------------------------} --- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs. --- --- > fromList [] == empty --- > fromList [(5,"a"), (3,"b"), (5, "c")] == fromList [(5,"c"), (3,"b")] --- > fromList [(5,"c"), (3,"b"), (5, "a")] == fromList [(5,"a"), (3,"b")] - -fromList :: [(Key,a)] -> IntMap a -fromList xs - = Foldable.foldl' ins empty xs - where - ins t (k,x) = insert k x t - --- | /O(n*min(n,W))/. Create a map from a list of key\/value pairs with a combining function. See also 'fromAscListWith'. --- --- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] --- > fromListWith (++) [] == empty - -fromListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWith f xs - = fromListWithKey (\_ x y -> f x y) xs - --- | /O(n*min(n,W))/. Build a map from a list of key\/value pairs with a combining function. See also fromAscListWithKey'. --- --- > fromListWith (++) [(5,"a"), (5,"b"), (3,"b"), (3,"a"), (5,"a")] == fromList [(3, "ab"), (5, "aba")] --- > fromListWith (++) [] == empty - -fromListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromListWithKey f xs - = 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. --- --- > 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. --- /The precondition (input list is ascending) is not checked./ --- --- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] - -fromAscListWith :: (a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWith f = 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. --- /The precondition (input list is ascending) is not checked./ --- --- > fromAscListWith (++) [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "ba")] - -fromAscListWithKey :: (Key -> a -> a -> a) -> [(Key,a)] -> IntMap a -fromAscListWithKey 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. --- /The precondition (input list is strictly ascending) is not checked./ --- --- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")] - -fromDistinctAscList :: [(Key,a)] -> IntMap a -fromDistinctAscList = 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. - -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 diff --git a/containers/src/Utils/Containers/Internal/IsList.hs b/containers/src/Utils/Containers/Internal/IsList.hs new file mode 100644 index 000000000..01b0e4dc9 --- /dev/null +++ b/containers/src/Utils/Containers/Internal/IsList.hs @@ -0,0 +1,14 @@ +{-# LANGUAGE CPP, Trustworthy #-} + +-- | A @Trustworthy@ module for the sole purpose of exposing 'GHC.Exts.IsList'. +-- To ensure that this is not misused, this module is not and should not be +-- exposed, and we only export the safe methods ('fromList') and ('toList'). +module Utils.Containers.Internal.IsList ( +#if __GLASGOW_HASKELL__ >= 708 + IsList(Item, fromList, toList) +#endif +) where + +#if __GLASGOW_HASKELL__ >= 708 +import GHC.Exts (IsList(Item, fromList, toList)) +#endif