Skip to content
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

Added HasCallStack to partial functions #493

Open
wants to merge 5 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
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
47 changes: 33 additions & 14 deletions Data/IntMap/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -331,6 +331,10 @@ import qualified Control.Category as Category
import Data.Coerce
#endif

#if __GLASGOW_HASKELL__ >= 800
import GHC.Stack (HasCallStack)
#endif


-- A "Nat" is a natural machine word (an unsigned Int)
type Nat = Word
Expand Down Expand Up @@ -390,8 +394,20 @@ bitmapOf x = shiftLL 1 (x .&. IntSet.suffixBitMask)
-- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
-- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'

#if __GLASGOW_HASKELL__ >= 800
(!) :: HasCallStack => IntMap a -> Key -> a
#else
(!) :: IntMap a -> Key -> a
(!) m k = find k m
#endif
(!) m0 !k = go m0
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))/. Find the value at a key.
-- Returns 'Nothing' when the element can not be found.
Expand Down Expand Up @@ -583,19 +599,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 Expand Up @@ -2169,11 +2172,19 @@ minView :: IntMap a -> Maybe (a, IntMap a)
minView t = liftM (first snd) (minViewWithKey t)

-- | /O(min(n,W))/. Delete and find the maximal element.
#if __GLASGOW_HASKELL__ >= 800
deleteFindMax :: HasCallStack => IntMap a -> ((Key, a), IntMap a)
#else
deleteFindMax :: IntMap a -> ((Key, a), IntMap a)
#endif
deleteFindMax = fromMaybe (error "deleteFindMax: empty map has no maximal element") . maxViewWithKey

-- | /O(min(n,W))/. Delete and find the minimal element.
#if __GLASGOW_HASKELL__ >= 800
deleteFindMin :: HasCallStack => IntMap a -> ((Key, a), IntMap a)
#else
deleteFindMin :: IntMap a -> ((Key, a), IntMap a)
#endif
deleteFindMin = fromMaybe (error "deleteFindMin: empty map has no minimal element") . minViewWithKey

-- | /O(min(n,W))/. The minimal key of the map. Returns 'Nothing' if the map is empty.
Expand All @@ -2188,7 +2199,11 @@ lookupMin (Bin _ m l r)
go Nil = Nothing

-- | /O(min(n,W))/. The minimal key of the map. Calls 'error' if the map is empty.
#if __GLASGOW_HASKELL__ >= 800
findMin :: HasCallStack => IntMap a -> (Key, a)
#else
findMin :: IntMap a -> (Key, a)
#endif
findMin t
| Just r <- lookupMin t = r
| otherwise = error "findMin: empty map has no minimal element"
Expand All @@ -2205,7 +2220,11 @@ lookupMax (Bin _ m l r)
go Nil = Nothing

-- | /O(min(n,W))/. The maximal key of the map. Calls 'error' if the map is empty.
#if __GLASGOW_HASKELL__ >= 800
findMax :: HasCallStack => IntMap a -> (Key, a)
#else
findMax :: IntMap a -> (Key, a)
#endif
findMax t
| Just r <- lookupMax t = r
| otherwise = error "findMax: empty map has no maximal element"
Expand Down
19 changes: 19 additions & 0 deletions Data/IntSet/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -216,6 +216,9 @@ import qualified GHC.Exts as GHCExts
import GHC.Prim (indexInt8OffAddr#)
#endif

#if __GLASGOW_HASKELL__ >= 800
import GHC.Stack (HasCallStack)
#endif

infixl 9 \\{-This comment teaches CPP correct behaviour -}

Expand Down Expand Up @@ -793,18 +796,30 @@ minView t =
-- | /O(min(n,W))/. Delete and find the minimal element.
--
-- > deleteFindMin set = (findMin set, deleteMin set)
#if __GLASGOW_HASKELL__ >= 800
deleteFindMin :: HasCallStack => IntSet -> (Key, IntSet)
#else
deleteFindMin :: IntSet -> (Key, IntSet)
#endif
deleteFindMin = fromMaybe (error "deleteFindMin: empty set has no minimal element") . minView

-- | /O(min(n,W))/. Delete and find the maximal element.
--
-- > deleteFindMax set = (findMax set, deleteMax set)
#if __GLASGOW_HASKELL__ >= 800
deleteFindMax :: HasCallStack => IntSet -> (Key, IntSet)
#else
deleteFindMax :: IntSet -> (Key, IntSet)
#endif
deleteFindMax = fromMaybe (error "deleteFindMax: empty set has no maximal element") . maxView


-- | /O(min(n,W))/. The minimal element of the set.
#if __GLASGOW_HASKELL__ >= 800
findMin :: HasCallStack => IntSet -> Key
#else
findMin :: IntSet -> Key
#endif
findMin Nil = error "findMin: empty set has no minimal element"
Copy link
Contributor

Choose a reason for hiding this comment

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

Does the call stack grow through the recursion here or elsewhere? If so, that's a big problem. You can check the Core to be sure. We know there are no Nils except at the root, but GHC does not! If the stacks build, you'll need to restructure the functions to fix that. Watch out for performance. If the times for the current benchmarks exercising these functions are too short to trust, consider adding more benchmarks.

Copy link
Author

Choose a reason for hiding this comment

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

Wow, yes it does. I didn't realize it would do that. The fix is to use an internal go function. I'll update the PR with that in a moment.

On a related note, I realized that my HasCallStack for Map.! was wrong because Map.! calls find which actually throws the error. Is there a reason it uses find rather than lookup? My inclination is to remove find altogether (it's only used by ! and looks identical other than the Maybe wrapper) and replace with lookup---then ! can call error itself, which should then make HasCallStack work the way we want.

Copy link
Author

Choose a reason for hiding this comment

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

Oh, I just noticed the [Note: Local 'go' functions and capturing] -- I'll read that and try to make sure I'm not doing anything stupid.

findMin (Tip kx bm) = kx + lowestBitSet bm
findMin (Bin _ m l r)
Expand All @@ -815,7 +830,11 @@ findMin (Bin _ m l r)
find Nil = error "findMin Nil"

-- | /O(min(n,W))/. The maximal element of a set.
#if __GLASGOW_HASKELL__ >= 800
findMax :: HasCallStack => IntSet -> Key
#else
findMax :: IntSet -> Key
#endif
findMax Nil = error "findMax: empty set has no maximal element"
findMax (Tip kx bm) = kx + highestBitSet bm
findMax (Bin _ m l r)
Expand Down
94 changes: 77 additions & 17 deletions Data/Map/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -411,6 +411,9 @@ import qualified Control.Category as Category
import Data.Coerce
#endif

#if __GLASGOW_HASKELL__ >= 800
import GHC.Stack (HasCallStack)
#endif

{--------------------------------------------------------------------
Operators
Expand All @@ -423,8 +426,18 @@ infixl 9 !,!?,\\ --
-- > fromList [(5,'a'), (3,'b')] ! 1 Error: element not in the map
-- > fromList [(5,'a'), (3,'b')] ! 5 == 'a'

#if __GLASGOW_HASKELL__ >= 800
(!) :: (HasCallStack, Ord k) => Map k a -> k -> a
#else
(!) :: Ord k => Map k a -> k -> a
(!) m k = find k m
#endif
(!) m !k = go m
where
go Tip = error "Map.!: given key is not an element in the map"
go (Bin _ kx x l r) = case compare k kx of
LT -> go l
GT -> go r
EQ -> x
#if __GLASGOW_HASKELL__
{-# INLINE (!) #-}
#endif
Expand Down Expand Up @@ -602,22 +615,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 @@ -1433,7 +1430,11 @@ alterFYoneda = go
-- > findIndex 6 (fromList [(5,"a"), (3,"b")]) Error: element is not in the map

-- See Note: Type of local 'go' function
#if __GLASGOW_HASKELL__ >= 800
findIndex :: (HasCallStack, Ord k) => k -> Map k a -> Int
#else
findIndex :: Ord k => k -> Map k a -> Int
#endif
findIndex = go 0
where
go :: Ord k => Int -> k -> Map k a -> Int
Expand Down Expand Up @@ -1477,6 +1478,18 @@ lookupIndex = go 0
-- > elemAt 1 (fromList [(5,"a"), (3,"b")]) == (5, "a")
-- > elemAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range

#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)
elemAt !_ Tip = error "Map.elemAt: index out of range"
elemAt i (Bin _ kx x l r)
Expand All @@ -1486,6 +1499,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 @@ -1566,6 +1580,21 @@ splitAt i0 m0
-- > updateAt (\_ _ -> Nothing) 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
-- > updateAt (\_ _ -> Nothing) (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range

#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
updateAt f !i t =
case t of
Expand All @@ -1578,6 +1607,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 @@ -1588,6 +1618,19 @@ updateAt f !i t =
-- > deleteAt 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
-- > deleteAt (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range

#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
deleteAt :: Int -> Map k a -> Map k a
deleteAt !i t =
case t of
Expand All @@ -1598,6 +1641,7 @@ deleteAt !i t =
EQ -> glue l r
where
sizeL = size l
#endif


{--------------------------------------------------------------------
Expand All @@ -1624,7 +1668,11 @@ lookupMin (Bin _ k x l _) = Just $! lookupMinSure k x l
-- > findMin (fromList [(5,"a"), (3,"b")]) == (3,"b")
-- > findMin empty Error: empty map has no minimal element

#if __GLASGOW_HASKELL__ >= 800
findMin :: HasCallStack => Map k a -> (k,a)
#else
findMin :: Map k a -> (k,a)
#endif
findMin t
| Just r <- lookupMin t = r
| otherwise = error "Map.findMin: empty map has no minimal element"
Expand All @@ -1649,7 +1697,11 @@ lookupMax :: Map k a -> Maybe (k, a)
lookupMax Tip = Nothing
lookupMax (Bin _ k x _ r) = Just $! lookupMaxSure k x r

#if __GLASGOW_HASKELL__ >= 800
findMax :: HasCallStack => Map k a -> (k,a)
#else
findMax :: Map k a -> (k,a)
#endif
findMax t
| Just r <- lookupMax t = r
| otherwise = error "Map.findMax: empty map has no maximal element"
Expand Down Expand Up @@ -3866,7 +3918,11 @@ maxViewSure = go
-- > deleteFindMin (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((3,"b"), fromList[(5,"a"), (10,"c")])
-- > deleteFindMin Error: can not return the minimal element of an empty map

#if __GLASGOW_HASKELL__ >= 800
deleteFindMin :: HasCallStack => Map k a -> ((k,a),Map k a)
#else
deleteFindMin :: Map k a -> ((k,a),Map k a)
#endif
deleteFindMin t = case minViewWithKey t of
Nothing -> (error "Map.deleteFindMin: can not return the minimal element of an empty map", Tip)
Just res -> res
Expand All @@ -3876,7 +3932,11 @@ deleteFindMin t = case minViewWithKey t of
-- > deleteFindMax (fromList [(5,"a"), (3,"b"), (10,"c")]) == ((10,"c"), fromList [(3,"b"), (5,"a")])
-- > deleteFindMax empty Error: can not return the maximal element of an empty map

#if __GLASGOW_HASKELL__ >= 800
deleteFindMax :: HasCallStack => Map k a -> ((k,a),Map k a)
#else
deleteFindMax :: Map k a -> ((k,a),Map k a)
#endif
deleteFindMax t = case maxViewWithKey t of
Nothing -> (error "Map.deleteFindMax: can not return the maximal element of an empty map", Tip)
Just res -> res
Expand Down
20 changes: 20 additions & 0 deletions Data/Map/Strict/Internal.hs
Original file line number Diff line number Diff line change
Expand Up @@ -418,6 +418,10 @@ import Data.Coerce
import Data.Functor.Identity (Identity (..))
#endif

#if __GLASGOW_HASKELL__ >= 800
import GHC.Stack (HasCallStack)
#endif


-- $strictness
--
Expand Down Expand Up @@ -881,6 +885,21 @@ atKeyIdentity k f t = Identity $ atKeyPlain Strict k (coerce f) t
-- > updateAt (\_ _ -> Nothing) 2 (fromList [(5,"a"), (3,"b")]) Error: index out of range
-- > updateAt (\_ _ -> Nothing) (-1) (fromList [(5,"a"), (3,"b")]) Error: index out of range

#if __GLASGOW_HASKELL__ >= 800
updateAt :: HasCallStack => (k -> a -> Maybe a) -> Int -> Map k a -> Map k a
updateAt = go where
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
updateAt f i t = i `seq`
case t of
Expand All @@ -893,6 +912,7 @@ updateAt f i t = i `seq`
Nothing -> glue l r
where
sizeL = size l
#endif

{--------------------------------------------------------------------
Minimal, Maximal
Expand Down
Loading