-
Notifications
You must be signed in to change notification settings - Fork 178
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Optimize IntMap.alter using unboxed sums. #523
base: master
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -11,6 +11,10 @@ | |
#if __GLASGOW_HASKELL__ >= 708 | ||
{-# LANGUAGE TypeFamilies #-} | ||
#endif | ||
#if __GLASGOW_HASKELL__ >= 802 | ||
{-# LANGUAGE UnboxedSums #-} | ||
{-# LANGUAGE UnboxedTuples #-} | ||
#endif | ||
|
||
{-# OPTIONS_HADDOCK not-home #-} | ||
|
||
|
@@ -311,6 +315,9 @@ import Prelude hiding (lookup, map, filter, foldr, foldl, null) | |
import Data.IntSet.Internal (Key) | ||
import qualified Data.IntSet.Internal as IntSet | ||
import Utils.Containers.Internal.BitUtil | ||
#if __GLASGOW_HASKELL__ >= 802 | ||
import Utils.Containers.Internal.PtrEquality (ptrEq) | ||
#endif | ||
import Utils.Containers.Internal.StrictFold | ||
import Utils.Containers.Internal.StrictPair | ||
|
||
|
@@ -937,6 +944,48 @@ updateLookupWithKey _ _ Nil = (Nothing,Nil) | |
-- | /O(min(n,W))/. The expression (@'alter' f k map@) alters the value @x@ at @k@, or absence thereof. | ||
-- 'alter' can be used to insert, delete, or update a value in an 'IntMap'. | ||
-- In short : @'lookup' k ('alter' f k m) = f ('lookup' k m)@. | ||
|
||
#if __GLASGOW_HASKELL__ >= 802 | ||
alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a | ||
alter f !k t = case alter# f k t of | ||
(# (# #) | #) -> t | ||
(# | t' #) -> t' | ||
{-# INLINE alter #-} | ||
|
||
|
||
-- Internal implementation which keeps track of whether or not the intmap was | ||
-- modified using an unboxed sum (Maybe). | ||
-- | ||
-- If no modifications are made to the map (# (# #) | #) is returned, otherwise | ||
-- (# | newMap #) is returned. | ||
alter# :: (Maybe a -> Maybe a) -> Key -> IntMap a -> (# (# #) | IntMap a #) | ||
alter# f !k t@(Bin p m l r) | ||
| nomatch k p m = case f Nothing of | ||
Nothing -> (# (# #) | #) | ||
Just x -> (# | link k (Tip k x) p t #) | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't think you're being strict enough with your return values. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do you mean something like the following? let !t' = link k (Tip k x) p t in (# | t' #) Or something else? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Yes, but see the link below for a clearer way. |
||
| zero k m = case alter# f k l of | ||
(# (# #) | #) -> (# (# #) | #) | ||
(# | l' #) -> (# | binCheckLeft p m l' r #) | ||
|
||
| otherwise = case alter# f k r of | ||
(# (# #) | #) -> (# (# #) | #) | ||
(# | r' #) -> (# | binCheckRight p m l r' #) | ||
|
||
alter# f k t@(Tip ky y) | ||
| k==ky = case f (Just y) of | ||
Just x -> if x `ptrEq` y | ||
then (# (# #) | #) | ||
else (# | Tip ky x #) | ||
Nothing -> (# | Nil #) | ||
| otherwise = case f Nothing of | ||
Just x -> (# | link k (Tip k x) ky t #) | ||
Nothing -> (# (# #) | #) | ||
alter# f k Nil = case f Nothing of | ||
Just x -> (# | Tip k x #) | ||
Nothing -> (# (# #) | #) | ||
{-# INLINABLE alter# #-} | ||
|
||
#else | ||
alter :: (Maybe a -> Maybe a) -> Key -> IntMap a -> IntMap a | ||
alter f !k t@(Bin p m l r) | ||
| nomatch k p m = case f Nothing of | ||
|
@@ -954,6 +1003,7 @@ alter f k t@(Tip ky y) | |
alter f k Nil = case f Nothing of | ||
Just x -> Tip k x | ||
Nothing -> Nil | ||
#endif | ||
|
||
-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ at | ||
-- @k@, or absence thereof. 'alterF' can be used to inspect, insert, delete, | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
What happens if you make the function have type
Maybe# a -> Maybe# a
, whereMaybe#
is the unboxed version ofMaybe
? Couldn't we often avoid theMaybe
allocation that way? Benchmark!There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
That function is user defined (passed to
alter
) so at some point you would have to do the conversion fromMaybe
toMaybe#
. As far as I know coerce wouldn't work forf
because they have different representations (Maybe
andMaybe#
that is), but there may be another way that I'm unaware of. Did you have a specific approach in mind?There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think it's a decent bet that the function we're passed will be small enough to inline. So calling
alter#
with(essentially) should usually avoid any actual
Maybe
s. Or so I imagine.