From 366a181fa26993fa3d4a12f3a80196113f9a93bf Mon Sep 17 00:00:00 2001
From: David Feuer <David.Feuer@gmail.com>
Date: Sat, 19 Nov 2022 17:49:31 -0500
Subject: [PATCH] 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.
---
 containers/containers.cabal                   |   1 +
 containers/src/Data/Map/Internal.hs           | 144 +++++++++++++++---
 containers/src/Data/Map/Strict/Internal.hs    |  44 +++---
 .../Utils/Containers/Internal/UnboxedMaybe.hs |  39 +++++
 4 files changed, 183 insertions(+), 45 deletions(-)
 create mode 100644 containers/src/Utils/Containers/Internal/UnboxedMaybe.hs

diff --git a/containers/containers.cabal b/containers/containers.cabal
index 30ac269e6..b39942add 100644
--- a/containers/containers.cabal
+++ b/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
diff --git a/containers/src/Data/Map/Internal.hs b/containers/src/Data/Map/Internal.hs
index 31275ac3b..a1c191e35 100644
--- a/containers/src/Data/Map/Internal.hs
+++ b/containers/src/Data/Map/Internal.hs
@@ -3,6 +3,8 @@
 {-# LANGUAGE PatternGuards #-}
 #if defined(__GLASGOW_HASKELL__)
 {-# LANGUAGE DeriveLift #-}
+{-# LANGUAGE MagicHash #-}
+{-# LANGUAGE UnboxedTuples #-}
 {-# LANGUAGE RoleAnnotations #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE Trustworthy #-}
@@ -10,10 +12,6 @@
 #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,6 +3006,7 @@ 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.
 --
@@ -2984,10 +3014,52 @@ mapMaybe f = mapMaybeWithKey (\_ x -> f x)
 -- > 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
diff --git a/containers/src/Data/Map/Strict/Internal.hs b/containers/src/Data/Map/Strict/Internal.hs
index 7b9b2b795..279237c74 100644
--- a/containers/src/Data/Map/Strict/Internal.hs
+++ b/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,6 +1274,7 @@ 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.
 --
@@ -1278,10 +1282,18 @@ mapMaybe f = mapMaybeWithKey (\_ x -> f x)
 -- > 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)\).
diff --git a/containers/src/Utils/Containers/Internal/UnboxedMaybe.hs b/containers/src/Utils/Containers/Internal/UnboxedMaybe.hs
new file mode 100644
index 000000000..4360d435f
--- /dev/null
+++ b/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# #-}