Skip to content

Added HasCallStack to partial functions #493

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

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 3 additions & 14 deletions Data/IntMap/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -399,7 +399,9 @@ bitmapOf x = shiftLL 1 (x .&. IntSet.suffixBitMask)
#else
(!) :: IntMap a -> Key -> a
#endif
(!) m k = find k m
(!) m k
| Just a <- lookup k m = a
| otherwise = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Using lookup will incur a Just allocation. One possible solution (for very recent GHC versions only) would be to define a lookup# function producing an unboxed sum:

type Maybe# a = (# Void# | a #)

pattern Nothing# :: Maybe# a
pattern Nothing# <- (# _ | #) where
  Nothing# = (# void# | #)

pattern Just# :: a -> Maybe# a
pattern Just# a = (# | a #)

lookup# :: Key -> IntMap a -> Maybe# a
lookup# k m = ...

lookup k m = case lookup# k m of
  Nothing# -> Nothing
  Just# a -> Just a

The (potentially) great thing about this is that the Just gets applied on the "outside", so inlining and the case-of-case transformation will end up making it go away altogether:

m ! k = case
    (case lookup# k m of
       Nothing# -> Nothing
       Just# a -> Just a) of
       Nothing -> error ...
       Just v -> v

-- ==> case-of-case, case of known constructor

m ! k =
  case lookup# k m of
       Nothing# -> error ...
       Just# a -> a

Unlike a CPS version, this theoretically helps user-written code as well. However, I have not actually tried benchmarking anything like this. And the CPP is potentially troublesome too.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This seems cool -- it seems like it could be a fairly big change that would even prevent an allocation when one uses lookup (or related functions that return Maybe values) and immediately unwraps -- but it's a bit big for what I was hoping to contribute just now (also, a bit over my head). I think I'll take the easy way out and go back to using a find function to get a result without the extra allocation. Perhaps my next PR will be to overhaul the backend to use unboxed sums where possible :).


-- | /O(min(n,W))/. Find the value at a key.
-- Returns 'Nothing' when the element can not be found.
Expand Down Expand Up @@ -591,19 +593,6 @@ lookup !k = go
go Nil = Nothing


-- See Note: Local 'go' functions and capturing]
find :: Key -> IntMap a -> a
find !k = go
where
go (Bin p m l r) | nomatch k p m = not_found
| zero k m = go l
| otherwise = go r
go (Tip kx x) | k == kx = x
| otherwise = not_found
go Nil = not_found

not_found = error ("IntMap.!: key " ++ show k ++ " is not an element of the map")

-- | /O(min(n,W))/. The expression @('findWithDefault' def k map)@
-- returns the value at key @k@ or returns @def@ when the key is not an
-- element of the map.
Expand Down
57 changes: 37 additions & 20 deletions Data/Map/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -431,7 +431,9 @@ infixl 9 !,!?,\\ --
#else
(!) :: Ord k => Map k a -> k -> a
#endif
(!) m k = find k m
(!) m k
| Just a <- lookup k m = a
| otherwise = error "Map.!: given key is not an element in the map"
#if __GLASGOW_HASKELL__
{-# INLINE (!) #-}
#endif
Expand Down Expand Up @@ -609,22 +611,6 @@ notMember k m = not $ member k m
{-# INLINE notMember #-}
#endif

-- | /O(log n)/. Find the value at a key.
-- Calls 'error' when the element can not be found.
find :: Ord k => k -> Map k a -> a
find = go
where
go !_ Tip = error "Map.!: given key is not an element in the map"
go k (Bin _ kx x l r) = case compare k kx of
LT -> go k l
GT -> go k r
EQ -> x
#if __GLASGOW_HASKELL__
{-# INLINABLE find #-}
#else
{-# INLINE find #-}
#endif

-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns
-- the value at key @k@ or returns default value @def@
-- when the key is not in the map.
Expand Down Expand Up @@ -1490,9 +1476,17 @@ lookupIndex = go 0

#if __GLASGOW_HASKELL__ >= 800
elemAt :: HasCallStack => Int -> Map k a -> (k,a)
elemAt = go where
go !_ Tip = error "Map.elemAt: index out of range"
go i (Bin _ kx x l r)
= case compare i sizeL of
LT -> elemAt i l
GT -> elemAt (i-sizeL-1) r
EQ -> (kx,x)
where
sizeL = size l
#else
elemAt :: Int -> Map k a -> (k,a)
#endif
elemAt !_ Tip = error "Map.elemAt: index out of range"
elemAt i (Bin _ kx x l r)
= case compare i sizeL of
Expand All @@ -1501,6 +1495,7 @@ elemAt i (Bin _ kx x l r)
EQ -> (kx,x)
where
sizeL = size l
#endif

-- | Take a given number of entries in key order, beginning
-- with the smallest keys.
Expand Down Expand Up @@ -1583,9 +1578,20 @@ splitAt i0 m0

#if __GLASGOW_HASKELL__ >= 800
updateAt :: HasCallStack => (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
updateAt = go where
go f !i t =
case t of
Tip -> error "Map.updateAt: index out of range"
Bin sx kx x l r -> case compare i sizeL of
LT -> balanceR kx x (go f i l) r
GT -> balanceL kx x l (go f (i-sizeL-1) r)
EQ -> case f kx x of
Just x' -> Bin sx kx x' l r
Nothing -> glue l r
where
sizeL = size l
#else
updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
#endif
updateAt f !i t =
case t of
Tip -> error "Map.updateAt: index out of range"
Expand All @@ -1597,6 +1603,7 @@ updateAt f !i t =
Nothing -> glue l r
where
sizeL = size l
#endif

-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in
-- the sequence sorted by keys. If the /index/ is out of range (less than zero,
Expand All @@ -1609,9 +1616,18 @@ updateAt f !i t =

#if __GLASGOW_HASKELL__ >= 800
deleteAt :: HasCallStack => Int -> Map k a -> Map k a
deleteAt = go where
go !i t =
case t of
Tip -> error "Map.deleteAt: index out of range"
Bin _ kx x l r -> case compare i sizeL of
LT -> balanceR kx x (go i l) r
GT -> balanceL kx x l (go (i-sizeL-1) r)
EQ -> glue l r
where
sizeL = size l
#else
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same question here.

deleteAt :: Int -> Map k a -> Map k a
#endif
deleteAt !i t =
case t of
Tip -> error "Map.deleteAt: index out of range"
Expand All @@ -1621,6 +1637,7 @@ deleteAt !i t =
EQ -> glue l r
where
sizeL = size l
#endif


{--------------------------------------------------------------------
Expand Down
14 changes: 13 additions & 1 deletion Data/Map/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -887,9 +887,20 @@ atKeyIdentity k f t = Identity $ atKeyPlain Strict k (coerce f) t

#if __GLASGOW_HASKELL__ >= 800
updateAt :: HasCallStack => (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
updateAt = go where
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this actually give the right call stacks? Thinking about this more carefully, I suspect you want to check if the operation will succeed, and only then call a go function that takes no call stack. The alternative would be to call the go function using withFrozenCallStack, which I suspect will be more expensive. Benchmarks may tell.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My tests show that, by default, functions like go do not have the HasCallStack constraint unless it is explicit in their type signatures. Therefore, this seems to give the right call stack. There's no need to check if the operation succeeds or not first, and because go doesn't have the constraint, there's no need for withFrozenCallStack.

go f i t = i `seq`
case t of
Tip -> error "Map.updateAt: index out of range"
Bin sx kx x l r -> case compare i sizeL of
LT -> balanceR kx x (go f i l) r
GT -> balanceL kx x l (go f (i-sizeL-1) r)
EQ -> case f kx x of
Just x' -> x' `seq` Bin sx kx x' l r
Nothing -> glue l r
where
sizeL = size l
#else
updateAt :: (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
#endif
updateAt f i t = i `seq`
case t of
Tip -> error "Map.updateAt: index out of range"
Expand All @@ -901,6 +912,7 @@ updateAt f i t = i `seq`
Nothing -> glue l r
where
sizeL = size l
#endif

{--------------------------------------------------------------------
Minimal, Maximal
Expand Down
4 changes: 0 additions & 4 deletions Data/Sequence/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -4505,11 +4505,7 @@ unstableSortBy cmp (Seq xs) =
-- | fromList2, given a list and its length, constructs a completely
-- balanced Seq whose elements are that list using the replicateA
-- generalization.
#if __GLASGOW_HASKELL__ >= 800
fromList2 :: HasCallStack => Int -> [a] -> Seq a
#else
fromList2 :: Int -> [a] -> Seq a
#endif
fromList2 n = execState (replicateA n (State ht))
where
ht (x:xs) = (xs, x)
Expand Down
23 changes: 21 additions & 2 deletions Data/Set/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -1254,9 +1254,17 @@ lookupIndex = go 0

#if __GLASGOW_HASKELL__ >= 800
elemAt :: HasCallStack => Int -> Set a -> a
elemAt = go where
go !_ Tip = error "Set.elemAt: index out of range"
go i (Bin _ x l r)
= case compare i sizeL of
LT -> go i l
GT -> go (i-sizeL-1) r
EQ -> x
where
sizeL = size l
#else
elemAt :: Int -> Set a -> a
#endif
elemAt !_ Tip = error "Set.elemAt: index out of range"
elemAt i (Bin _ x l r)
= case compare i sizeL of
Expand All @@ -1265,6 +1273,7 @@ elemAt i (Bin _ x l r)
EQ -> x
where
sizeL = size l
#endif

-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based index in
-- the sorted sequence of elements. If the /index/ is out of range (less than zero,
Expand All @@ -1279,9 +1288,18 @@ elemAt i (Bin _ x l r)

#if __GLASGOW_HASKELL__ >= 800
deleteAt :: HasCallStack => Int -> Set a -> Set a
deleteAt = go where
go !i t =
case t of
Tip -> error "Set.deleteAt: index out of range"
Bin _ x l r -> case compare i sizeL of
LT -> balanceR x (go i l) r
GT -> balanceL x l (go (i-sizeL-1) r)
EQ -> glue l r
where
sizeL = size l
#else
deleteAt :: Int -> Set a -> Set a
#endif
deleteAt !i t =
case t of
Tip -> error "Set.deleteAt: index out of range"
Expand All @@ -1291,6 +1309,7 @@ deleteAt !i t =
EQ -> glue l r
where
sizeL = size l
#endif

-- | Take a given number of elements in order, beginning
-- with the smallest ones.
Expand Down