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# #-}