Improve min and max view laziness and inlining
authorDavid Feuer <David.Feuer@gmail.com>
Mon, 8 Jan 2018 21:26:44 +0000 (16:26 -0500)
committerDavid Feuer <David.Feuer@gmail.com>
Mon, 22 Jan 2018 06:16:11 +0000 (01:16 -0500)
* Harmonize laziness (not strictness!) of min and max views between
`IntMap` and `Map`.

* Improve GHC's ability to unbox the results of min and max views
  for `IntMap`.

Data/IntMap/Internal.hs
Data/Map/Internal.hs

index 9b30673..a2ff125 100644 (file)
@@ -301,7 +301,6 @@ import Data.Functor.Classes
 #endif
 
 import Control.DeepSeq (NFData(rnf))
-import Control.Monad (liftM)
 import Data.Bits
 import qualified Data.Foldable as Foldable
 import Data.Maybe (fromMaybe)
@@ -2107,17 +2106,26 @@ data View a = View {-# UNPACK #-} !Key a !(IntMap a)
 -- > maxViewWithKey empty == Nothing
 
 maxViewWithKey :: IntMap a -> Maybe ((Key, a), IntMap a)
-maxViewWithKey t =
+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 -> Nothing
+    Nil -> error "maxViewWithKeySure Nil"
     Bin p m l r | m < 0 ->
-      Just $ case go l of View k a l' -> ((k, a), binCheckLeft p m l' r)
-    _ -> Just $ case go t of View k a t' -> ((k, a), t')
+      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 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.
@@ -2129,14 +2137,31 @@ 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 ->
-      Just $ case go r of View k a r' -> ((k, a), binCheckRight p m l r')
-    _ -> Just $ case go t of View k a t' -> ((k, a), t')
+      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 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.
 --
@@ -2154,25 +2179,25 @@ updateMax f = updateMaxWithKey (const f)
 updateMin :: (a -> Maybe a) -> IntMap a -> IntMap a
 updateMin f = updateMinWithKey (const f)
 
--- Similar to the Arrow instance.
-first :: (a -> c) -> (a, b) -> (c, b)
-first f (x,y) = (f x,y)
-
 -- | /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 = liftM (first snd) (maxViewWithKey t)
+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 = liftM (first snd) (minViewWithKey t)
+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
 
@@ -2188,6 +2213,7 @@ lookupMin (Bin _ m l r)
           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
@@ -2205,6 +2231,7 @@ lookupMax (Bin _ m l 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
index e35e043..e0f7592 100644 (file)
@@ -1725,9 +1725,13 @@ updateMaxWithKey f (Bin _ kx x l r)    = balanceL kx x l (updateMaxWithKey f r)
 
 minViewWithKey :: Map k a -> Maybe ((k,a), Map k a)
 minViewWithKey Tip = Nothing
-minViewWithKey (Bin _ k x l r) =
+minViewWithKey (Bin _ k x l r) = Just $
   case minViewSure k x l r of
-    MinView km xm t -> Just ((km, xm), t)
+    MinView km xm t -> ((km, xm), t)
+-- We inline this to give GHC the best possible chance of getting
+-- rid of the Maybe and pair constructors, as well as the thunk under
+-- the Just.
+{-# INLINE minViewWithKey #-}
 
 -- | /O(log n)/. Retrieves the maximal (key,value) pair of the map, and
 -- the map stripped of that element, or 'Nothing' if passed an empty map.
@@ -1737,9 +1741,11 @@ minViewWithKey (Bin _ k x l r) =
 
 maxViewWithKey :: Map k a -> Maybe ((k,a), Map k a)
 maxViewWithKey Tip = Nothing
-maxViewWithKey (Bin _ k x l r) =
+maxViewWithKey (Bin _ k x l r) = Just $
   case maxViewSure k x l r of
-    MaxView km xm t -> Just ((km, xm), t)
+    MaxView km xm t -> ((km, xm), t)
+-- See note on inlining at minViewWithKey
+{-# INLINE maxViewWithKey #-}
 
 -- | /O(log n)/. Retrieves the value associated with minimal key of the
 -- map, and the map stripped of that element, or 'Nothing' if passed an
@@ -1751,7 +1757,7 @@ maxViewWithKey (Bin _ k x l r) =
 minView :: Map k a -> Maybe (a, Map k a)
 minView t = case minViewWithKey t of
               Nothing -> Nothing
-              Just ((_, x), t') -> Just (x, t')
+              Just ~((_, x), t') -> Just (x, t')
 
 -- | /O(log n)/. Retrieves the value associated with maximal key of the
 -- map, and the map stripped of that element, or 'Nothing' if passed an
@@ -1763,7 +1769,7 @@ minView t = case minViewWithKey t of
 maxView :: Map k a -> Maybe (a, Map k a)
 maxView t = case maxViewWithKey t of
               Nothing -> Nothing
-              Just ((_, x), t') -> Just (x, t')
+              Just ~((_, x), t') -> Just (x, t')
 
 {--------------------------------------------------------------------
   Union.
@@ -3852,6 +3858,7 @@ minViewSure = go
     go k x (Bin _ kl xl ll lr) r =
       case go kl xl ll lr of
         MinView km xm l' -> MinView km xm (balanceR k x l' r)
+{-# NOINLINE minViewSure #-}
 
 maxViewSure :: k -> a -> Map k a -> Map k a -> MaxView k a
 maxViewSure = go
@@ -3860,6 +3867,7 @@ maxViewSure = go
     go k x l (Bin _ kr xr rl rr) =
       case go kr xr rl rr of
         MaxView km xm r' -> MaxView km xm (balanceL k x l r')
+{-# NOINLINE maxViewSure #-}
 
 -- | /O(log n)/. Delete and find the minimal element.
 --