Adding HasCallStack to Partial functions

containers/src/Data/IntMap/Internal.hs

@@ -364,6 +364,9 @@ import Text.Read
 #endif
 import qualified Control.Category as Category
 
+#if __GLASGOW_HASKELL__ >= 800
+import GHC.Stack (HasCallStack)
+#endif
 
 {--------------------------------------------------------------------
   Types
@@ -429,8 +432,20 @@ deriving instance Lift a => Lift (IntMap a)
 -- > 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 l r) | left k p  = 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.
@@ -648,18 +663,6 @@ lookup !k = go
                   | otherwise = Nothing
     go Nil = Nothing
 
--- See Note: Local 'go' functions and capturing]
-find :: Key -> IntMap a -> a
-find !k = go
-  where
-    go (Bin p l r) | left k p  = 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.
@@ -2351,15 +2354,23 @@ minView t = fmap (\((_, x), t') -> (x, t')) (minViewWithKey t)
 -- Calls 'error' if the map is empty.
 --
 -- __Note__: This function is partial. Prefer 'maxViewWithKey'.
+#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.
 --
 -- Calls 'error' if the map is empty.
 --
 -- __Note__: This function is partial. Prefer 'minViewWithKey'.
+#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
 
 -- The KeyValue type is used when returning a key-value pair and helps with
@@ -2394,7 +2405,11 @@ lookupMin (Bin p l r) =
 -- | \(O(\min(n,W))\). The minimal key of the map. Calls 'error' if the map is empty.
 --
 -- __Note__: This function is partial. Prefer 'lookupMin'.
+#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"
@@ -2415,7 +2430,11 @@ lookupMax (Bin p l r) =
 -- | \(O(\min(n,W))\). The maximal key of the map. Calls 'error' if the map is empty.
 --
 -- __Note__: This function is partial. Prefer 'lookupMax'.
+#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"

containers/src/Data/IntSet/Internal.hs

@@ -234,6 +234,10 @@ import Language.Haskell.TH ()
 import qualified Data.Foldable as Foldable
 import Data.Functor.Identity (Identity(..))
 
+#if __GLASGOW_HASKELL__ >= 800
+import GHC.Stack (HasCallStack)
+#endif
+
 infixl 9 \\{-This comment teaches CPP correct behaviour -}
 
 {--------------------------------------------------------------------
@@ -1103,15 +1107,23 @@ minView t =
 -- Calls 'error' if the set is empty.
 --
 -- __Note__: This function is partial. Prefer 'minView'.
+#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.
 --
 -- Calls 'error' if the set is empty.
 --
 -- __Note__: This function is partial. Prefer 'maxView'.
+#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
 
 lookupMinSure :: IntSet -> Key
@@ -1133,7 +1145,11 @@ lookupMin (Bin p l r) = Just $! lookupMinSure (if signBranch p then r else l)
 -- is empty.
 --
 -- __Note__: This function is partial. Prefer 'lookupMin'.
+#if __GLASGOW_HASKELL__ >= 800
+findMin :: HasCallStack => IntSet -> Key
+#else
 findMin :: IntSet -> Key
+#endif
 findMin t
   | Just r <- lookupMin t = r
   | otherwise = error "findMin: empty set has no minimal element"
@@ -1157,7 +1173,11 @@ lookupMax (Bin p l r) = Just $! lookupMaxSure (if signBranch p then l else r)
 -- is empty.
 --
 -- __Note__: This function is partial. Prefer 'lookupMax'.
+#if __GLASGOW_HASKELL__ >= 800
+findMax :: HasCallStack => IntSet -> Key
+#else
 findMax :: IntSet -> Key
+#endif
 findMax t
   | Just r <- lookupMax t = r
   | otherwise = error "findMax: empty set has no maximal element"

containers/src/Data/Map/Internal.hs

@@ -426,6 +426,9 @@ import Data.Coerce
 import Text.Read hiding (lift)
 #endif
 import qualified Control.Category as Category
+#if __GLASGOW_HASKELL__ >= 800
+import GHC.Stack (HasCallStack)
+#endif
 
 {--------------------------------------------------------------------
   Operators
@@ -440,8 +443,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
@@ -626,20 +639,6 @@ notMember k m = not $ member k m
 {-# INLINE notMember #-}
 #endif
 
-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.
@@ -1484,7 +1483,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
@@ -1530,6 +1533,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)
@@ -1539,6 +1554,7 @@ elemAt i (Bin _ kx x l r)
       EQ -> (kx,x)
   where
     sizeL = size l
+#endif
 
 -- | \(O(\log n)\). Take a given number of entries in key order, beginning
 -- with the smallest keys.
@@ -1621,6 +1637,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
@@ -1633,6 +1664,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,
@@ -1645,6 +1677,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
@@ -1655,6 +1700,7 @@ deleteAt !i t =
       EQ -> glue l r
       where
         sizeL = size l
+#endif
 
 
 {--------------------------------------------------------------------
@@ -1702,7 +1748,11 @@ lookupMin (Bin _ k x l _) = Just $! kvToTuple (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"
@@ -1730,7 +1780,11 @@ lookupMax (Bin _ k x _ r) = Just $! kvToTuple (lookupMaxSure k x r)
 -- > findMax (fromList [(5,"a"), (3,"b")]) == (5,"a")
 -- > findMax empty                            Error: empty map has no maximal element
 
+#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"
@@ -4182,7 +4236,11 @@ maxViewSure !k x !l r = case r of
 -- Calls 'error' if the map is empty.
 --
 -- __Note__: This function is partial. Prefer 'minViewWithKey'.
+#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
@@ -4192,7 +4250,11 @@ deleteFindMin t = case minViewWithKey t of
 -- Calls 'error' if the map is empty.
 --
 -- __Note__: This function is partial. Prefer 'maxViewWithKey'.
+#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

containers/src/Data/Map/Strict/Internal.hs

@@ -424,6 +424,10 @@ import Data.Coerce
 import Data.Functor.Identity (Identity (..))
 #endif
 
+#if __GLASGOW_HASKELL__ >= 800
+import GHC.Stack (HasCallStack)
+#endif
+
 import qualified Data.Foldable as Foldable
 
 -- [Note: Pointer equality for sharing]
@@ -869,6 +873,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
@@ -881,6 +900,7 @@ updateAt f i t = i `seq`
               Nothing -> glue l r
       where
         sizeL = size l
+#endif
 
 {--------------------------------------------------------------------
   Minimal, Maximal