Stop using hedge algorithms
[packages/containers.git] / Data / Map / Strict.hs
index 766ca11..a837304 100644 (file)
@@ -173,10 +173,16 @@ module Data.Map.Strict
     , fromAscListWith
     , fromAscListWithKey
     , fromDistinctAscList
+    , fromDescList
+    , fromDescListWith
+    , fromDescListWithKey
+    , fromDistinctDescList
 
     -- * Filter
     , filter
     , filterWithKey
+    , restrictKeys
+    , withoutKeys
     , partition
     , partitionWithKey
 
@@ -268,6 +274,10 @@ import Data.Map.Base hiding
     , fromAscListWith
     , fromAscListWithKey
     , fromDistinctAscList
+    , fromDescList
+    , fromDescListWith
+    , fromDescListWithKey
+    , fromDistinctDescList
     , mapMaybe
     , mapMaybeWithKey
     , mapEither
@@ -652,12 +662,13 @@ alter = go
 --
 -- Note on rewrite rules:
 --
--- This module includes GHC rewrite rules to optimize 'alterF' for the 'Const',
--- 'Identity', and @(,) b@ functors. In general, these rules improve
--- performance. The main exception is that when using 'Identity', deleting a
--- key that is already absent takes longer than it would without the rules. If
--- you expect this to occur a very large fraction of the time, you might
--- consider using a private copy of the 'Identity' type.
+-- This module includes GHC rewrite rules to optimize 'alterF' for
+-- the 'Const' and 'Identity' functors. In general, these rules
+-- improve performance. The sole exception is that when using
+-- 'Identity', deleting a key that is already absent takes longer
+-- than it would without the rules. If you expect this to occur
+-- a very large fraction of the time, you might consider using a
+-- private copy of the 'Identity' type.
 --
 -- Note: 'alterF' is a flipped version of the 'at' combinator from
 -- 'Control.Lens.At'.
@@ -674,13 +685,7 @@ alterF f k m = atKeyImpl Strict k f m
 -- `Control.Applicative.Const` and just doing a lookup.
 {-# RULES
 "alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m
-"alterF/Pair" forall k (f :: Maybe a -> (b, Maybe a)) . alterF f k = atKeyPair k f
  #-}
-
-atKeyPair :: Ord k => k -> (Maybe a -> (b, Maybe a)) -> Map k a -> (b, Map k a)
-atKeyPair k f t = atKeyWithLookup Strict k f t
-{-# INLINABLE atKeyPair #-}
-
 #if MIN_VERSION_base(4,8,0)
 -- base 4.8 and above include Data.Functor.Identity, so we can
 -- save a pretty decent amount of time by handling it specially.
@@ -790,19 +795,18 @@ unionsWith f ts
 {--------------------------------------------------------------------
   Union with a combining function
 --------------------------------------------------------------------}
--- | /O(n+m)/. Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
+-- | /O(m*log(n/m + 1)), m <= n/. Union with a combining function.
 --
 -- > unionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "aA"), (7, "C")]
 
 unionWith :: Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
-unionWith f m1 m2
-  = unionWithKey (\_ x y -> f x y) m1 m2
+unionWith f t1 t2 = mergeWithKey (\_ x1 x2 -> Just $ f x1 x2) id id t1 t2
 #if __GLASGOW_HASKELL__
 {-# INLINABLE unionWith #-}
 #endif
 
--- | /O(n+m)/.
--- Union with a combining function. The implementation uses the efficient /hedge-union/ algorithm.
+-- | /O(m*log(n/m + 1)), m <= n/.
+-- Union with a combining function.
 --
 -- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value
 -- > unionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == fromList [(3, "b"), (5, "5:a|A"), (7, "C")]
@@ -822,15 +826,13 @@ unionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) id id t1 t2
 -- encountered, the combining function is applied to the values of these keys.
 -- If it returns 'Nothing', the element is discarded (proper set difference). If
 -- it returns (@'Just' y@), the element is updated with a new value @y@.
--- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 --
 -- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing
 -- > differenceWith f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (7, "C")])
 -- >     == singleton 3 "b:B"
 
 differenceWith :: Ord k => (a -> b -> Maybe a) -> Map k a -> Map k b -> Map k a
-differenceWith f m1 m2
-  = differenceWithKey (\_ x y -> f x y) m1 m2
+differenceWith f t1 t2 = mergeWithKey (\_ x1 x2 -> f x1 x2) id (const Tip) t1 t2
 #if __GLASGOW_HASKELL__
 {-# INLINABLE differenceWith #-}
 #endif
@@ -839,7 +841,6 @@ differenceWith f m1 m2
 -- encountered, the combining function is applied to the key and both values.
 -- If it returns 'Nothing', the element is discarded (proper set difference). If
 -- it returns (@'Just' y@), the element is updated with a new value @y@.
--- The implementation uses an efficient /hedge/ algorithm comparable with /hedge-union/.
 --
 -- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing
 -- > differenceWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (3, "B"), (10, "C")])
@@ -856,25 +857,21 @@ differenceWithKey f t1 t2 = mergeWithKey f id (const Tip) t1 t2
   Intersection
 --------------------------------------------------------------------}
 
--- | /O(n+m)/. Intersection with a combining function.  The implementation uses
--- an efficient /hedge/ algorithm comparable with /hedge-union/.
+-- | /O(m*log(n/m + 1)), m <= n/. Intersection with a combining function.
 --
 -- > intersectionWith (++) (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "aA"
 
 intersectionWith :: Ord k => (a -> b -> c) -> Map k a -> Map k b -> Map k c
-intersectionWith f m1 m2
-  = intersectionWithKey (\_ x y -> f x y) m1 m2
+intersectionWith f t1 t2 = mergeWithKey (\_ x1 x2 -> Just $ f x1 x2) (const Tip) (const Tip) t1 t2
 #if __GLASGOW_HASKELL__
 {-# INLINABLE intersectionWith #-}
 #endif
 
--- | /O(n+m)/. Intersection with a combining function.  The implementation uses
--- an efficient /hedge/ algorithm comparable with /hedge-union/.
+-- | /O(m*log(n/m + 1)), m <= n/. Intersection with a combining function.
 --
 -- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar
 -- > intersectionWithKey f (fromList [(5, "a"), (3, "b")]) (fromList [(5, "A"), (7, "C")]) == singleton 5 "5:a|A"
 
-
 intersectionWithKey :: Ord k => (k -> a -> b -> c) -> Map k a -> Map k b -> Map k c
 intersectionWithKey f t1 t2 = mergeWithKey (\k x1 x2 -> Just $ f k x1 x2) (const Tip) (const Tip) t1 t2
 #if __GLASGOW_HASKELL__
@@ -928,22 +925,19 @@ mergeWithKey f g1 g2 = go
   where
     go Tip t2 = g2 t2
     go t1 Tip = g1 t1
-    go t1 t2 = hedgeMerge NothingS NothingS t1 t2
-
-    hedgeMerge _   _   t1  Tip = g1 t1
-    hedgeMerge blo bhi Tip (Bin _ kx x l r) = g2 $ link kx x (filterGt blo l) (filterLt bhi r)
-    hedgeMerge blo bhi (Bin _ kx x l r) t2 = let l' = hedgeMerge blo bmi l (trim blo bmi t2)
-                                                 (found, trim_t2) = trimLookupLo kx bhi t2
-                                                 r' = hedgeMerge bmi bhi r trim_t2
-                                             in case found of
-                                                  Nothing -> case g1 (singleton kx x) of
-                                                               Tip -> merge l' r'
-                                                               (Bin _ _ x' Tip Tip) -> link kx x' l' r'
-                                                               _ -> error "mergeWithKey: Given function only1 does not fulfil required conditions (see documentation)"
-                                                  Just x2 -> case f kx x x2 of
-                                                               Nothing -> merge l' r'
-                                                               Just x' -> x' `seq` link kx x' l' r'
-      where bmi = JustS kx
+    go (Bin _ kx x l1 r1) t2 =
+      case found of
+        Nothing -> case g1 (singleton kx x) of
+                     Tip -> merge l' r'
+                     (Bin _ _ x' Tip Tip) -> link kx x' l' r'
+                     _ -> error "mergeWithKey: Given function only1 does not fulfill required conditions (see documentation)"
+        Just x2 -> case f kx x x2 of
+                     Nothing -> merge l' r'
+                     Just x' -> link kx x' l' r'
+      where
+        (l2, found, r2) = splitLookup kx t2
+        l' = go l1 l2
+        r' = go r1 r2
 {-# INLINE mergeWithKey #-}
 
 {--------------------------------------------------------------------
@@ -1010,8 +1004,13 @@ mapEitherWithKey f0 t0 = toPair $ go f0 t0
 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
 
 map :: (a -> b) -> Map k a -> Map k b
-map _ Tip = Tip
-map f (Bin sx kx x l r) = let x' = f x in x' `seq` Bin sx kx x' (map f l) (map f r)
+map f = go
+  where
+    go Tip = Tip
+    go (Bin sx kx x l r) = let !x' = f x in Bin sx kx x' (go l) (go r)
+-- We use `go` to let `map` inline. This is important if `f` is a constant
+-- function.
+
 #ifdef __GLASGOW_HASKELL__
 {-# NOINLINE [1] map #-}
 {-# RULES
@@ -1220,10 +1219,15 @@ fromListWithKey f xs
 {--------------------------------------------------------------------
   Building trees from ascending/descending lists can be done in linear time.
 
-  Note that if [xs] is ascending that:
+  Note that if [xs] is ascending then:
     fromAscList xs       == fromList xs
     fromAscListWith f xs == fromListWith f xs
+
+  If [xs] is descending then:
+    fromDescList xs       == fromList xs
+    fromDescListWith f xs == fromListWith f xs
 --------------------------------------------------------------------}
+
 -- | /O(n)/. Build a map from an ascending list in linear time.
 -- /The precondition (input list is ascending) is not checked./
 --
@@ -1231,7 +1235,6 @@ fromListWithKey f xs
 -- > fromAscList [(3,"b"), (5,"a"), (5,"b")] == fromList [(3, "b"), (5, "b")]
 -- > valid (fromAscList [(3,"b"), (5,"a"), (5,"b")]) == True
 -- > valid (fromAscList [(5,"a"), (3,"b"), (5,"b")]) == False
-
 fromAscList :: Eq k => [(k,a)] -> Map k a
 fromAscList xs
   = fromAscListWithKey (\_ x _ -> x) xs
@@ -1239,6 +1242,20 @@ fromAscList xs
 {-# INLINABLE fromAscList #-}
 #endif
 
+-- | /O(n)/. Build a map from a descending list in linear time.
+-- /The precondition (input list is descending) is not checked./
+--
+-- > fromDescList [(5,"a"), (3,"b")]          == fromList [(3, "b"), (5, "a")]
+-- > fromDescList [(5,"a"), (5,"b"), (3,"a")] == fromList [(3, "b"), (5, "b")]
+-- > valid (fromDescList [(5,"a"), (5,"b"), (3,"b")]) == True
+-- > valid (fromDescList [(5,"a"), (3,"b"), (5,"b")]) == False
+fromDescList :: Eq k => [(k,a)] -> Map k a
+fromDescList xs
+  = fromDescListWithKey (\_ x _ -> x) xs
+#if __GLASGOW_HASKELL__
+{-# INLINABLE fromDescList #-}
+#endif
+
 -- | /O(n)/. Build a map from an ascending list in linear time with a combining function for equal keys.
 -- /The precondition (input list is ascending) is not checked./
 --
@@ -1253,6 +1270,20 @@ fromAscListWith f xs
 {-# INLINABLE fromAscListWith #-}
 #endif
 
+-- | /O(n)/. Build a map from a descending list in linear time with a combining function for equal keys.
+-- /The precondition (input list is descending) is not checked./
+--
+-- > fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "ba")]
+-- > valid (fromDescListWith (++) [(5,"a"), (5,"b"), (3,"b")]) == True
+-- > valid (fromDescListWith (++) [(5,"a"), (3,"b"), (5,"b")]) == False
+
+fromDescListWith :: Eq k => (a -> a -> a) -> [(k,a)] -> Map k a
+fromDescListWith f xs
+  = fromDescListWithKey (\_ x y -> f x y) xs
+#if __GLASGOW_HASKELL__
+{-# INLINABLE fromDescListWith #-}
+#endif
+
 -- | /O(n)/. Build a map from an ascending list in linear time with a
 -- combining function for equal keys.
 -- /The precondition (input list is ascending) is not checked./
@@ -1281,6 +1312,34 @@ fromAscListWithKey f xs
 {-# INLINABLE fromAscListWithKey #-}
 #endif
 
+-- | /O(n)/. Build a map from a descending list in linear time with a
+-- combining function for equal keys.
+-- /The precondition (input list is descending) is not checked./
+--
+-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2
+-- > fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")] == fromList [(3, "b"), (5, "5:b5:ba")]
+-- > valid (fromDescListWithKey f [(5,"a"), (5,"b"), (5,"b"), (3,"b")]) == True
+-- > valid (fromDescListWithKey f [(5,"a"), (3,"b"), (5,"b"), (5,"b")]) == False
+
+fromDescListWithKey :: Eq k => (k -> a -> a -> a) -> [(k,a)] -> Map k a
+fromDescListWithKey f xs
+  = fromDistinctDescList (combineEq f xs)
+  where
+  -- [combineEq f xs] combines equal elements with function [f] in an ordered list [xs]
+  combineEq _ xs'
+    = case xs' of
+        []     -> []
+        [x]    -> [x]
+        (x:xx) -> combineEq' x xx
+
+  combineEq' z [] = [z]
+  combineEq' z@(kz,zz) (x@(kx,xx):xs')
+    | kx==kz    = let yy = f kx xx zz in yy `seq` combineEq' (kx,yy) xs'
+    | otherwise = z:combineEq' x xs'
+#if __GLASGOW_HASKELL__
+{-# INLINABLE fromDescListWithKey #-}
+#endif
+
 -- | /O(n)/. Build a map from an ascending list of distinct elements in linear time.
 -- /The precondition is not checked./
 --
@@ -1305,3 +1364,28 @@ fromDistinctAscList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip T
                       res@(_, []) -> res
                       (l, (ky, y):ys) -> case create (s `shiftR` 1) ys of
                         (r, zs) -> y `seq` (link ky y l r, zs)
+
+-- | /O(n)/. Build a map from a descending list of distinct elements in linear time.
+-- /The precondition is not checked./
+--
+-- > fromDistinctDescList [(5,"a"), (3,"b")] == fromList [(3, "b"), (5, "a")]
+-- > valid (fromDistinctDescList [(5,"a"), (3,"b")])          == True
+-- > valid (fromDistinctDescList [(5,"a"), (3,"b"), (3,"a")]) == False
+
+-- For some reason, when 'singleton' is used in fromDistinctDescList or in
+-- create, it is not inlined, so we inline it manually.
+fromDistinctDescList :: [(k,a)] -> Map k a
+fromDistinctDescList [] = Tip
+fromDistinctDescList ((kx0, x0) : xs0) = x0 `seq` go (1::Int) (Bin 1 kx0 x0 Tip Tip) xs0
+  where
+    go !_ t [] = t
+    go s r ((kx, x) : xs) = case create s xs of
+                              (l, ys) -> x `seq` go (s `shiftL` 1) (link kx x l r) ys
+
+    create !_ [] = (Tip, [])
+    create s xs@(x' : xs')
+      | s == 1 = case x' of (kx, x) -> x `seq` (Bin 1 kx x Tip Tip, xs')
+      | otherwise = case create (s `shiftR` 1) xs of
+                      res@(_, []) -> res
+                      (r, (ky, y):ys) -> case create (s `shiftR` 1) ys of
+                        (l, zs) -> y `seq` (link ky y l r, zs)