Speed up IntMap
authorDavid Feuer <David.Feuer@gmail.com>
Sun, 8 May 2016 00:54:24 +0000 (20:54 -0400)
committerDavid Feuer <David.Feuer@gmail.com>
Sun, 8 May 2016 19:29:06 +0000 (15:29 -0400)
`delete`, `alter`, `update`, etc., used a `bin` smart
constructor to avoid installing any non-root `Nil`s. Now only
the ones that could have become `Nil` are checked, which is
a good bit cheaper since they're in cache. `adjustWithKey`
was implemented using `updateWithKey`, but in fact it never
needs to worry about `Nil`s, so implementing it directly
eliminates all such checks.

Make `updateLookupWithKey` in `Data.IntMap.Lazy` strict in its
recursive call to avoid essentially useless lazy pair allocation.

Data/IntMap/Base.hs
Data/IntMap/Strict.hs

index b0994e6..e0a462c 100644 (file)
@@ -207,6 +207,8 @@ module Data.IntMap.Base (
     , intFromNat
     , link
     , bin
+    , binCheckLeft
+    , binCheckRight
     , zero
     , nomatch
     , match
@@ -717,8 +719,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 +746,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 +781,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 +803,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 +822,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
@@ -1052,10 +1061,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 +1077,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 +1095,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 +1111,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"
 
@@ -2100,6 +2109,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
index b919033..64fdd9d 100644 (file)
@@ -439,8 +439,16 @@ 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 =
+  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
@@ -469,8 +477,8 @@ updateWithKey f !k t =
   case t of
     Bin p m l r
       | nomatch k p m -> t
-      | zero k m      -> bin p m (updateWithKey f k l) r
-      | otherwise     -> bin p m l (updateWithKey f k r)
+      | 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'
@@ -495,8 +503,8 @@ updateLookupWithKey f0 !k0 t0 = toPair $ go f0 k0 t0
       case t of
         Bin p m l r
           | nomatch k p m -> (Nothing :*: t)
-          | zero k m      -> let (found :*: l') = go f k l in (found :*: bin p m l' r)
-          | otherwise     -> let (found :*: r') = go f k r in (found :*: bin p m l r')
+          | 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')
@@ -516,8 +524,8 @@ alter f !k t =
       | 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)
     Tip ky y
       | k==ky         -> case f (Just y) of
                            Just !x -> Tip ky x
@@ -667,10 +675,10 @@ mergeWithKey f g1 g2 = mergeWithKey' bin combine g1 g2
 
 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
@@ -683,10 +691,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