Unboxing and streamlining Map maps

* Use an unboxed-sum version of `Maybe` to implement `mapMaybeWithKey`.
  This potentially (I suspect usually) allows all the `Maybe`s to be
  erased.

* Comprehensive rewrite rules for both strict and lazy versions of
  `map`, `mapWithKey`, `mapMaybeWithKey`, and `filterWithKey` quickly
  get out of hand. Following `unordered-containers`, tame the mess
  by implementing both lazy and strict mapping functions in terms of
  versions that use unboxed results. Rewrite rules on these underlying
  functions will then apply uniformly. One concern: I found it a bit
  tricky to get the unfoldings I wanted; lots of things had to be marked
  `INLINABLE` explicitly.

Author: treeowl

containers/containers.cabal

@@ -69,6 +69,7 @@ Library
         Utils.Containers.Internal.BitUtil
         Utils.Containers.Internal.BitQueue
         Utils.Containers.Internal.StrictPair
+        Utils.Containers.Internal.UnboxedMaybe
 
     other-modules:
         Prelude

containers/src/Data/Map/Internal.hs

@@ -3,17 +3,15 @@
 {-# LANGUAGE PatternGuards #-}
 #if defined(__GLASGOW_HASKELL__)
 {-# LANGUAGE DeriveLift #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE RoleAnnotations #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE Trustworthy #-}
 {-# LANGUAGE TypeFamilies #-}
 #endif
 #define USE_MAGIC_PROXY 1
 
-#ifdef USE_MAGIC_PROXY
-{-# LANGUAGE MagicHash #-}
-#endif
-
 {-# OPTIONS_HADDOCK not-home #-}
 
 #include "containers.h"
@@ -236,7 +234,9 @@ module Data.Map.Internal (
     -- * Traversal
     -- ** Map
     , map
+    , map#
     , mapWithKey
+    , mapWithKey#
     , traverseWithKey
     , traverseMaybeWithKey
     , mapAccum
@@ -301,6 +301,7 @@ module Data.Map.Internal (
 
     , mapMaybe
     , mapMaybeWithKey
+    , mapMaybeWithKey#
     , mapEither
     , mapEitherWithKey
 
@@ -407,6 +408,7 @@ import Data.Data
 import qualified Control.Category as Category
 import Data.Coerce
 #endif
+import Utils.Containers.Internal.UnboxedMaybe
 
 
 {--------------------------------------------------------------------
@@ -2849,6 +2851,7 @@ isProperSubmapOfBy f t1 t2
 filter :: (a -> Bool) -> Map k a -> Map k a
 filter p m
   = filterWithKey (\_ x -> p x) m
+{-# INLINE filter #-}
 
 -- | \(O(n)\). Filter all keys\/values that satisfy the predicate.
 --
@@ -2863,6 +2866,32 @@ filterWithKey p t@(Bin _ kx x l r)
   | otherwise = link2 pl pr
   where !pl = filterWithKey p l
         !pr = filterWithKey p r
+{-# NOINLINE [1] filterWithKey #-}
+
+{-# RULES
+"filterWK/filterWK" forall p q m. filterWithKey p (filterWithKey q m) =
+  filterWithKey (\k x -> q k x && p k x) m
+"filterWK/map#" forall p f m. filterWithKey p (map# f m) =
+  mapMaybeWithKey# (\k x -> case f x of
+    (# y #)
+      | p k y -> Just# y
+      | otherwise -> Nothing#) m
+"filterWK/mapWK#" forall p f m. filterWithKey p (mapWithKey# f m) =
+  mapMaybeWithKey# (\k x -> case f k x of
+    (# y #)
+      | p k y -> Just# y
+      | otherwise -> Nothing#) m
+"map#/filterWK" forall f p m. map# f (filterWithKey p m) =
+  mapMaybeWithKey# (\k x ->
+    if p k x
+    then case f x of (# y #) -> Just# y
+    else Nothing#) m
+"mapWK#/filterWK" forall f p m. mapWithKey# f (filterWithKey p m) =
+  mapMaybeWithKey# (\k x ->
+    if p k x
+    then case f k x of (# y #) -> Just# y
+    else Nothing#) m
+ #-}
 
 -- | \(O(n)\). Filter keys and values using an 'Applicative'
 -- predicate.
@@ -2977,17 +3006,60 @@ partitionWithKey p0 t0 = toPair $ go p0 t0
 
 mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b
 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
+{-# INLINE mapMaybe #-}
 
 -- | \(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 :: (k -> a -> Maybe b) -> Map k a -> Map k b
+{-
 mapMaybeWithKey _ Tip = Tip
 mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
   Just y  -> link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
   Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+-}
+mapMaybeWithKey f = \m ->
+  mapMaybeWithKey# (\k x -> toMaybe# (f k x)) m
+{-# INLINE mapMaybeWithKey #-}
+
+mapMaybeWithKey# :: (k -> a -> Maybe# b) -> Map k a -> Map k b
+mapMaybeWithKey# _ Tip = Tip
+mapMaybeWithKey# f (Bin _ kx x l r) = case f kx x of
+  Just# y  -> link kx y (mapMaybeWithKey# f l) (mapMaybeWithKey# f r)
+  Nothing# -> link2 (mapMaybeWithKey# f l) (mapMaybeWithKey# f r)
+{-# NOINLINE [1] mapMaybeWithKey# #-}
+
+{-# RULES
+"mapMaybeWK#/map#" forall f g m. mapMaybeWithKey# f (map# g m) =
+  mapMaybeWithKey# (\k x -> case g x of (# y #) -> f k y) m
+"map#/mapMaybeWK#" forall f g m. map# f (mapMaybeWithKey# g m) =
+  mapMaybeWithKey#
+    (\k x -> case g k x of
+               Nothing# -> Nothing#
+               Just# y -> case f y of (# z #) -> Just# z) m
+"mapMaybeWK#/mapWK#" forall f g m. mapMaybeWithKey# f (mapWithKey# g m) =
+  mapMaybeWithKey# (\k x -> case g k x of (# y #) -> f k y) m
+"mapWK#/mapMaybeWK#" forall f g m. mapWithKey# f (mapMaybeWithKey# g m) =
+  mapMaybeWithKey#
+    (\k x -> case g k x of
+               Nothing# -> Nothing#
+               Just# y -> case f k y of (# z #) -> Just# z) m
+"mapMaybeWK#/mapMaybeWK#" forall f g m. mapMaybeWithKey# f (mapMaybeWithKey# g m) =
+  mapMaybeWithKey#
+    (\k x -> case g k x of
+               Nothing# -> Nothing#
+               Just# y -> f k y) m
+"mapMaybeWK#/filterWK" forall f p m. mapMaybeWithKey# f (filterWithKey p m) =
+  mapMaybeWithKey# (\k x -> if p k x then f k x else Nothing#) m
+"filterWK/mapMaybeWK#" forall p f m. filterWithKey p (mapMaybeWithKey# f m) =
+  mapMaybeWithKey# (\k x -> case f k x of
+    Nothing# -> Nothing#
+    Just# y
+      | p k y -> Just# y
+      | otherwise -> Nothing#) m
+ #-}
 
 -- | \(O(n)\). Traverse keys\/values and collect the 'Just' results.
 --
@@ -3045,18 +3117,34 @@ 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
+#ifdef __GLASGOW_HASKELL__
+-- We define map using map# solely to reduce the number of rewrite
+-- rules we need.
+map f = map# (\x -> (# f x #))
+{-# INLINABLE map #-}
+#else
 map f = go where
   go Tip = Tip
   go (Bin sx kx x l r) = Bin sx kx (f x) (go l) (go r)
--- We use a `go` function to allow `map` to inline. This makes
--- a big difference if someone uses `map (const x) m` instead
--- of `x <$ m`; it doesn't seem to do any harm.
+#endif
 
 #ifdef __GLASGOW_HASKELL__
-{-# NOINLINE [1] map #-}
+map# :: (a -> (# b #)) -> Map k a -> Map k b
+map# f = go where
+  go Tip = Tip
+  go (Bin sx kx x l r)
+    | (# y #) <- f x
+    = Bin sx kx y (go l) (go r)
+-- We use a `go` function to allow `map#` to inline. Without this,
+-- we'd slow down both strict and lazy map, which wouldn't be great.
+-- This also lets us avoid a custom implementation of <$
+
+{-# NOINLINE [1] map# #-}
+-- Perhaps surprisingly, this map#/coerce rule seems to work. Hopefully,
+-- it will continue to do so.
 {-# RULES
-"map/map" forall f g xs . map f (map g xs) = map (f . g) xs
-"map/coerce" map coerce = coerce
+"map#/map#" forall f g xs . map# f (map# g xs) = map# (\x -> case g x of (# y #) -> f y) xs
+"map#/coerce" map# (\x -> (# coerce x #)) = coerce
  #-}
 #endif
 
@@ -3066,21 +3154,33 @@ map f = go where
 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
 
 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
+#ifdef __GLASGOW_HASKELL__
+mapWithKey f = mapWithKey# (\k a -> (# f k a #))
+{-# INLINABLE mapWithKey #-}
+#else
 mapWithKey _ Tip = Tip
 mapWithKey f (Bin sx kx x l r) = Bin sx kx (f kx x) (mapWithKey f l) (mapWithKey f r)
+#endif
+
+-- | A version of 'mapWithKey' that takes a function producing a unary
+-- unboxed tuple.
+mapWithKey# :: (k -> a -> (# b #)) -> Map k a -> Map k b
+mapWithKey# f = go where
+  go Tip = Tip
+  go (Bin sx kx x l r)
+    | (# y #) <- f kx x
+    = Bin sx kx y (go l) (go r)
 
 #ifdef __GLASGOW_HASKELL__
-{-# NOINLINE [1] mapWithKey #-}
+{-# 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
+"mapWK#/mapWK#" forall f g xs. mapWithKey# f (mapWithKey# g xs) = mapWithKey# (\k x -> case g k x of (# y #) -> f k y) xs
+"mapWK#/map#" forall f g xs. mapWithKey# f (map# g xs) = mapWithKey# (\k x -> case g x of (# y #) -> f k y) xs
+"map#/mapWK#" forall f g xs. map# f (mapWithKey# g xs) = mapWithKey# (\k x -> case g k x of (# y #) -> f y) xs
  #-}
 #endif
 
+
 -- | \(O(n)\).
 -- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@
 -- That is, behaves exactly like a regular 'traverse' except that the traversing
@@ -4195,10 +4295,12 @@ instance (Ord k, Read k) => Read1 (Map k) where
 --------------------------------------------------------------------}
 instance Functor (Map k) where
   fmap f m  = map f m
-#ifdef __GLASGOW_HASKELL__
-  _ <$ Tip = Tip
-  a <$ (Bin sx kx _ l r) = Bin sx kx a (a <$ l) (a <$ r)
-#endif
+  {-# INLINABLE fmap #-}
+  a <$ m = map (const a) m
+  -- For some reason, we need an explicit INLINE or INLINABLE pragma to
+  -- get the unfolding to use map rather than expanding into a recursive
+  -- function that RULES will never match. Hmm....
+  {-# INLINABLE (<$) #-}
 
 -- | Traverses in order of increasing key.
 instance Traversable (Map k) where

containers/src/Data/Map/Strict/Internal.hs

@@ -1,7 +1,9 @@
 {-# LANGUAGE CPP #-}
 {-# LANGUAGE BangPatterns #-}
 #if defined(__GLASGOW_HASKELL__)
+{-# LANGUAGE MagicHash #-}
 {-# LANGUAGE Trustworthy #-}
+{-# LANGUAGE UnboxedTuples #-}
 #endif
 {-# OPTIONS_HADDOCK not-home #-}
 
@@ -420,6 +422,7 @@ import Data.Semigroup (Arg (..))
 import qualified Data.Set.Internal as Set
 import qualified Data.Map.Internal as L
 import Utils.Containers.Internal.StrictPair
+import Utils.Containers.Internal.UnboxedMaybe (Maybe# (..))
 
 import Data.Bits (shiftL, shiftR)
 #ifdef __GLASGOW_HASKELL__
@@ -1271,17 +1274,26 @@ mergeWithKey f g1 g2 = go
 
 mapMaybe :: (a -> Maybe b) -> Map k a -> Map k b
 mapMaybe f = mapMaybeWithKey (\_ x -> f x)
+{-# INLINABLE mapMaybe #-}
 
 -- | \(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 :: (k -> a -> Maybe b) -> Map k a -> Map k b
+{-
+ -
 mapMaybeWithKey _ Tip = Tip
 mapMaybeWithKey f (Bin _ kx x l r) = case f kx x of
   Just y  -> y `seq` link kx y (mapMaybeWithKey f l) (mapMaybeWithKey f r)
   Nothing -> link2 (mapMaybeWithKey f l) (mapMaybeWithKey f r)
+-}
+mapMaybeWithKey f = \m ->
+  L.mapMaybeWithKey# (\k x -> case f k x of
+    Nothing -> Nothing#
+    Just !a -> Just# a) m
+{-# INLINABLE mapMaybeWithKey #-}
 
 -- | \(O(n)\). Traverse keys\/values and collect the 'Just' results.
 --
@@ -1340,19 +1352,16 @@ 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
+#ifdef __GLASGOW_HASKELL__
+map f = L.map# (\x -> let !y = f x in (# y #))
+{-# INLINABLE map #-}
+#else
 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
-"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.
@@ -1361,27 +1370,14 @@ map f = go
 -- > mapWithKey f (fromList [(5,"a"), (3,"b")]) == fromList [(3, "3:b"), (5, "5:a")]
 
 mapWithKey :: (k -> a -> b) -> Map k a -> Map k b
+#ifdef __GLASGOW_HASKELL__
+mapWithKey f = L.mapWithKey# (\k x -> let !y = f k x in (# y #))
+{-# INLINABLE mapWithKey #-}
+#else
 mapWithKey _ Tip = Tip
 mapWithKey f (Bin sx kx x l r) =
   let x' = f kx x
   in x' `seq` Bin sx kx x' (mapWithKey f l) (mapWithKey f r)
-
-#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/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)\).

containers/src/Utils/Containers/Internal/UnboxedMaybe.hs

@@ -0,0 +1,39 @@
+{-# language LambdaCase #-}
+{-# language MagicHash #-}
+{-# language PatternSynonyms #-}
+{-# language UnboxedSums #-}
+{-# language UnboxedTuples #-}
+{-# language UnliftedNewtypes #-}
+module Utils.Containers.Internal.UnboxedMaybe
+  ( Maybe# (Just#, Nothing#)
+  , maybe#
+  , toMaybe
+  , toMaybe#
+  ) where
+
+newtype Maybe# a = Maybe# (# (##) | a #)
+
+maybe# :: r -> (a -> r) -> Maybe# a -> r
+maybe# n j = \case
+  Nothing# -> n
+  Just# a -> j a
+{-# INLINE maybe# #-}
+
+pattern Nothing# :: Maybe# a
+pattern Nothing# = Maybe# (# (##)| #)
+{-# INLINE Nothing# #-}
+
+pattern Just# :: a -> Maybe# a
+pattern Just# a = Maybe# (#|a #)
+{-# INLINE Just# #-}
+
+{-# COMPLETE Just#, Nothing# #-}
+
+toMaybe :: Maybe# a -> Maybe a
+toMaybe = maybe# Nothing Just
+{-# INLINE toMaybe #-}
+
+toMaybe# :: Maybe a -> Maybe# a
+toMaybe# Nothing = Nothing#
+toMaybe# (Just a) = Just# a
+{-# INLINE toMaybe# #-}