Stop using hedge algorithms
[packages/containers.git] / Data / IntMap / Base.hs
index b0994e6..f45e7d1 100644 (file)
@@ -2,11 +2,11 @@
 {-# LANGUAGE BangPatterns #-}
 #if __GLASGOW_HASKELL__
 {-# LANGUAGE MagicHash, DeriveDataTypeable, StandaloneDeriving #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 #endif
 #if !defined(TESTING) && __GLASGOW_HASKELL__ >= 703
 {-# LANGUAGE Trustworthy #-}
 #endif
-{-# LANGUAGE ScopedTypeVariables #-}
 #if __GLASGOW_HASKELL__ >= 708
 {-# LANGUAGE TypeFamilies #-}
 #endif
@@ -90,6 +90,7 @@ module Data.IntMap.Base (
     , updateWithKey
     , updateLookupWithKey
     , alter
+    , alterF
 
     -- * Combine
 
@@ -163,6 +164,8 @@ module Data.IntMap.Base (
     -- * Filter
     , filter
     , filterWithKey
+    , restrictKeys
+    , withoutKeys
     , partition
     , partitionWithKey
 
@@ -207,6 +210,8 @@ module Data.IntMap.Base (
     , intFromNat
     , link
     , bin
+    , binCheckLeft
+    , binCheckRight
     , zero
     , nomatch
     , match
@@ -245,6 +250,9 @@ import Data.Utils.StrictPair
 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 ((<$))
+#endif
 #if __GLASGOW_HASKELL__ >= 708
 import qualified GHC.Exts as GHCExts
 #endif
@@ -717,8 +725,8 @@ insertLookupWithKey _ k x Nil = (Nothing,Tip k x)
 delete :: Key -> IntMap a -> IntMap a
 delete !k t@(Bin p m l r)
   | nomatch k p m = t
-  | zero k m      = bin p m (delete k l) r
-  | otherwise     = bin p m l (delete k r)
+  | 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
@@ -744,8 +752,15 @@ adjust f k m
 -- > adjustWithKey f 7 empty                         == empty
 
 adjustWithKey ::  (Key -> a -> a) -> Key -> IntMap a -> IntMap a
-adjustWithKey f
-  = updateWithKey (\k' x -> Just (f k' x))
+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
@@ -772,8 +787,8 @@ update f
 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      = bin p m (updateWithKey f k l) r
-  | otherwise     = bin p m l (updateWithKey f k r)
+  | 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'
@@ -794,8 +809,8 @@ updateWithKey _ _ Nil = Nil
 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,bin p m l' r)
-  | otherwise     = let (found,r') = updateLookupWithKey f k r in (found,bin p m l r')
+  | 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')
@@ -813,8 +828,8 @@ 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      = bin p m (alter f k l) r
-  | otherwise     = bin p m l (alter f k r)
+  | 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
@@ -826,6 +841,42 @@ 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:
+--
+-- @
+-- interactiveAlter :: Int -> IntMap String -> IO (IntMap String)
+-- interactiveAlter k m = alterF f k m where
+--   f Nothing -> do
+--      putStrLn $ show k ++
+--          " was not found in the map. Would you like to add it?"
+--      getUserResponse1 :: IO (Maybe String)
+--   f (Just old) -> do
+--      putStrLn "The key is currently bound to " ++ show old ++
+--          ". Would you like to change or delete it?"
+--      getUserresponse2 :: IO (Maybe String)
+-- @
+--
+-- 'alterF' is the most general operation for working with an individual
+-- key that may or may not be in a given map.
+--
+-- Note: 'alterF' is a flipped version of the 'at' combinator from
+-- 'Control.Lens.At'.
+--
+-- @since 0.5.8
+
+alterF :: Functor f
+       => (Maybe a -> f (Maybe a)) -> Key -> IntMap a -> f (IntMap a)
+-- This implementation was stolen from 'Control.Lens.At'.
+alterF f k m = (<$> f mv) $ \fres ->
+  case fres of
+    Nothing -> maybe m (const (delete k m)) mv
+    Just v' -> insert k v' m
+  where mv = lookup k m
 
 {--------------------------------------------------------------------
   Union
@@ -911,6 +962,49 @@ differenceWithKey :: (Key -> a -> b -> Maybe a) -> IntMap a -> IntMap b -> IntMa
 differenceWithKey f m1 m2
   = mergeWithKey f id (const Nil) m1 m2
 
+-- | 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 = go
+  where
+    go t1@(Bin p1 m1 l1 r1) t2@(IntSet.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      = t1
+      where
+        merge1 | nomatch p2 p1 m1  = t1
+               | zero p2 m1        = binCheckLeft p1 m1 (go l1 t2) r1
+               | otherwise         = binCheckRight p1 m1 l1 (go r1 t2)
+        merge2 | nomatch p1 p2 m2  = t1
+               | zero p1 m2        = bin p2 m2 (go t1 l2) Nil
+               | otherwise         = bin p2 m2 Nil (go t1 r2)
+
+    go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip k2' _) = merge t2' k2' t1'
+      where merge t2 k2 t1@(Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = t1
+                                             | zero k2 m1 = binCheckLeft p1 m1 (merge t2 k2 l1) r1
+                                             | otherwise  = binCheckRight p1 m1 l1 (merge t2 k2 r1)
+            merge _ k2 t1@(Tip k1 _) | k1 == k2 = Nil
+                                     | otherwise = t1
+            merge _ _  Nil = Nil
+
+    go t1@(Bin _ _ _ _) IntSet.Nil = t1
+
+    go t1'@(Tip k1' _) t2' = merge t1' k1' t2'
+      where merge t1 k1 (IntSet.Bin p2 m2 l2 r2) | nomatch k1 p2 m2 = t1
+                                                 | zero k1 m2 = bin p2 m2 (merge t1 k1 l2) Nil
+                                                 | otherwise  = bin p2 m2 Nil (merge t1 k1 r2)
+            merge t1 k1 (IntSet.Tip k2 _) | k1 == k2 = Nil
+                                          | otherwise = t1
+            merge t1 _  IntSet.Nil = t1
+
+    go Nil _ = Nil
+
 
 {--------------------------------------------------------------------
   Intersection
@@ -923,6 +1017,50 @@ intersection :: IntMap a -> IntMap b -> IntMap a
 intersection m1 m2
   = mergeWithKey' bin const (const Nil) (const Nil) m1 m2
 
+-- | /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 = go
+  where
+    go t1@(Bin p1 m1 l1 r1) t2@(IntSet.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      = Nil
+      where
+        merge1 | nomatch p2 p1 m1  = Nil
+               | zero p2 m1        = bin p1 m1 (go l1 t2) Nil
+               | otherwise         = bin p1 m1 Nil (go r1 t2)
+        merge2 | nomatch p1 p2 m2  = Nil
+               | zero p1 m2        = bin p2 m2 (go t1 l2) Nil
+               | otherwise         = bin p2 m2 Nil (go t1 r2)
+
+    go t1'@(Bin _ _ _ _) t2'@(IntSet.Tip k2' _) = merge t2' k2' t1'
+      where merge t2 k2 (Bin p1 m1 l1 r1) | nomatch k2 p1 m1 = Nil
+                                          | zero k2 m1 = bin p1 m1 (merge t2 k2 l1) Nil
+                                          | otherwise  = bin p1 m1 Nil (merge t2 k2 r1)
+            merge _ k2 t1@(Tip k1 _) | k1 == k2 = t1
+                                     | otherwise = Nil
+            merge _ _  Nil = Nil
+
+    go (Bin _ _ _ _) IntSet.Nil = Nil
+
+    go t1'@(Tip k1' _) t2' = merge t1' k1' t2'
+      where merge t1 k1 (IntSet.Bin p2 m2 l2 r2)
+              | nomatch k1 p2 m2 = Nil
+              | zero k1 m2 = bin p2 m2 (merge t1 k1 l2) Nil
+              | otherwise  = bin p2 m2 Nil (merge t1 k1 r2)
+            merge t1 k1 (IntSet.Tip k2 _) | k1 == k2 = t1
+                                          | otherwise = Nil
+            merge _ _  IntSet.Nil = Nil
+
+    go 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"
@@ -1052,10 +1190,10 @@ mergeWithKey' bin' f g1 g2 = go
 
 updateMinWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
 updateMinWithKey f t =
-  case t of Bin p m l r | m < 0 -> bin p m l (go f r)
+  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) = bin p m (go f' l) r
+    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
@@ -1068,10 +1206,10 @@ updateMinWithKey f t =
 
 updateMaxWithKey :: (Key -> a -> Maybe a) -> IntMap a -> IntMap a
 updateMaxWithKey f t =
-  case t of Bin p m l r | m < 0 -> bin p m (go f l) r
+  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) = bin p m l (go f' r)
+    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
@@ -1086,10 +1224,10 @@ updateMaxWithKey f t =
 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
 maxViewWithKey t =
   case t of Nil -> Nothing
-            Bin p m l r | m < 0 -> case go l of (result, l') -> Just (result, bin p m l' r)
+            Bin p m l r | m < 0 -> case go l of (result, l') -> Just (result, binCheckLeft p m l' r)
             _ -> Just (go t)
   where
-    go (Bin p m l r) = case go r of (result, r') -> (result, bin p m l r')
+    go (Bin p m l r) = case go r of (result, r') -> (result, binCheckRight p m l r')
     go (Tip k y) = ((k, y), Nil)
     go Nil = error "maxViewWithKey Nil"
 
@@ -1102,10 +1240,10 @@ maxViewWithKey t =
 minViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
 minViewWithKey t =
   case t of Nil -> Nothing
-            Bin p m l r | m < 0 -> case go r of (result, r') -> Just (result, bin p m l r')
+            Bin p m l r | m < 0 -> case go r of (result, r') -> Just (result, binCheckRight p m l r')
             _ -> Just (go t)
   where
-    go (Bin p m l r) = case go l of (result, l') -> (result, bin p m l' r)
+    go (Bin p m l r) = case go l of (result, l') -> (result, binCheckLeft p m l' r)
     go (Tip k y) = ((k, y), Nil)
     go Nil = error "minViewWithKey Nil"
 
@@ -1284,11 +1422,11 @@ isSubmapOfBy _         Nil _           = True
 -- > map (++ "x") (fromList [(5,"a"), (3,"b")]) == fromList [(3, "bx"), (5, "ax")]
 
 map :: (a -> b) -> IntMap a -> IntMap b
-map f t
-  = case t of
-      Bin p m l r -> Bin p m (map f l) (map f r)
-      Tip k x     -> Tip k (f x)
-      Nil         -> Nil
+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 #-}
@@ -1983,14 +2121,20 @@ fromAscListWithKey f (x0 : xs0) = fromDistinctAscList (combineEq x0 xs0)
 --
 -- > fromDistinctAscList [(3,"b"), (5,"a")] == fromList [(3, "b"), (5, "a")]
 
+#if __GLASGOW_HASKELL__
 fromDistinctAscList :: forall a. [(Key,a)] -> IntMap a
+#else
+fromDistinctAscList ::            [(Key,a)] -> IntMap a
+#endif
 fromDistinctAscList []         = Nil
 fromDistinctAscList (z0 : zs0) = work z0 zs0 Nada
   where
     work (kx,vx) []            stk = finish kx (Tip kx vx) stk
     work (kx,vx) (z@(kz,_):zs) stk = reduce z zs (branchMask kx kz) kx (Tip kx vx) stk
 
+#if __GLASGOW_HASKELL__
     reduce :: (Key,a) -> [(Key,a)] -> Mask -> Prefix -> IntMap a -> Stack a -> IntMap a
+#endif
     reduce z zs _ px tx Nada = work z zs (Push px tx Nada)
     reduce z zs m px tx stk@(Push py ty stk') =
         let mxy = branchMask px py
@@ -2044,6 +2188,12 @@ instance Ord a => Ord (IntMap a) where
 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
 --------------------------------------------------------------------}
@@ -2074,7 +2224,7 @@ instance (Read e) => Read (IntMap e) where
   Typeable
 --------------------------------------------------------------------}
 
-INSTANCE_TYPEABLE1(IntMap,intMapTc,"IntMap")
+INSTANCE_TYPEABLE1(IntMap)
 
 {--------------------------------------------------------------------
   Helpers
@@ -2100,6 +2250,17 @@ 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