Allow users to customise the collection of exception annotations

Add a global `CollectExceptionAnnotationMechanism` which determines how
`ExceptionAnnotation`s are collected upon throwing an `Exception`.

This API is exposed via `ghc-experimental`.

By overriding how we collect `Backtraces`, we can control how the
`Backtraces` are displayed to the user by newtyping `Backtraces` and
giving a different instance for `ExceptionAnnotation`.

A concrete use-case for this feature is allowing us to experiment with
alternative stack decoders, without having to modify `base`, which take
additional information from the stack frames.

This commit does not modify how `Backtraces` are currently
collected or displayed.

(cherry picked from commit dee28cd)

Author: fendor

libraries/ghc-experimental/src/GHC/Exception/Backtrace/Experimental.hs

@@ -0,0 +1,20 @@
+{-
+Module      :  GHC.Exception.Backtrace.Experimental
+Copyright   :  (c) The GHC Team
+License     :  see libraries/ghc-experimental/LICENSE
+
+Maintainer  :  [email protected]
+Stability   :  experimental
+Portability :  non-portable (GHC extensions)
+
+This module exposes experimental extensions to the Backtrace mechanism of GHC.
+-}
+module GHC.Exception.Backtrace.Experimental (
+  -- * Collecting exception annotations (like backtraces)
+  CollectExceptionAnnotationMechanism,
+  getCollectExceptionAnnotationMechanism,
+  setCollectExceptionAnnotation,
+  collectExceptionAnnotation,
+  ) where
+
+import GHC.Internal.Exception.Backtrace

libraries/ghc-internal/src/GHC/Internal/Exception.hs

@@ -70,7 +70,8 @@ import GHC.Internal.Show
 import GHC.Internal.Stack.Types
 import GHC.Internal.IO.Unsafe
 import {-# SOURCE #-} GHC.Internal.Stack (prettyCallStackLines, prettyCallStack, prettySrcLoc, withFrozenCallStack)
-import {-# SOURCE #-} GHC.Internal.Exception.Backtrace (collectBacktraces)
+import {-# SOURCE #-} GHC.Internal.Exception.Backtrace (collectExceptionAnnotation)
+import GHC.Internal.Exception.Context (SomeExceptionAnnotation(..))
 import GHC.Internal.Exception.Type
 
 -- | Throw an exception.  Exceptions may be thrown from purely
@@ -166,8 +167,8 @@ toExceptionWithBacktrace :: (HasCallStack, Exception e)
                          => e -> IO SomeException
 toExceptionWithBacktrace e
   | backtraceDesired e = do
-      bt <- collectBacktraces
-      return (addExceptionContext bt (toException e))
+      SomeExceptionAnnotation ea <- collectExceptionAnnotation
+      return (addExceptionContext ea (toException e))
   | otherwise = return (toException e)
 
 -- | This is thrown when the user calls 'error'. The @String@ is the

libraries/ghc-internal/src/GHC/Internal/Exception/Backtrace.hs

@@ -12,7 +12,7 @@ import GHC.Internal.IO.Unsafe (unsafePerformIO)
 import GHC.Internal.Exception.Context
 import GHC.Internal.Ptr
 import GHC.Internal.Data.Maybe (fromMaybe)
-import GHC.Internal.Stack.Types as GHC.Stack (CallStack)
+import GHC.Internal.Stack.Types as GHC.Stack (CallStack, HasCallStack)
 import qualified GHC.Internal.Stack as HCS
 import qualified GHC.Internal.ExecutionStack.Internal as ExecStack
 import qualified GHC.Internal.Stack.CloneStack as CloneStack
@@ -86,6 +86,37 @@ setBacktraceMechanismState bm enabled = do
     _ <- atomicModifyIORef'_ enabledBacktraceMechanismsRef (setBacktraceMechanismEnabled bm enabled)
     return ()
 
+-- | How to collect 'ExceptionAnnotation's on throwing 'Exception's.
+--
+data CollectExceptionAnnotationMechanism = CollectExceptionAnnotationMechanism
+  { ceaCollectExceptionAnnotationMechanism :: HasCallStack => IO SomeExceptionAnnotation
+  }
+
+defaultCollectExceptionAnnotationMechanism :: CollectExceptionAnnotationMechanism
+defaultCollectExceptionAnnotationMechanism = CollectExceptionAnnotationMechanism
+  { ceaCollectExceptionAnnotationMechanism = SomeExceptionAnnotation `fmap` collectBacktraces
+  }
+
+collectExceptionAnnotationMechanismRef :: IORef CollectExceptionAnnotationMechanism
+collectExceptionAnnotationMechanismRef =
+    unsafePerformIO $ newIORef defaultCollectExceptionAnnotationMechanism
+{-# NOINLINE collectExceptionAnnotationMechanismRef #-}
+
+-- | Returns the current callback for collecting 'ExceptionAnnotation's on throwing 'Exception's.
+--
+getCollectExceptionAnnotationMechanism :: IO CollectExceptionAnnotationMechanism
+getCollectExceptionAnnotationMechanism = readIORef collectExceptionAnnotationMechanismRef
+
+-- | Set the callback for collecting an 'ExceptionAnnotation'.
+--
+setCollectExceptionAnnotation :: ExceptionAnnotation a => (HasCallStack => IO a) -> IO ()
+setCollectExceptionAnnotation collector = do
+  let cea = CollectExceptionAnnotationMechanism
+        { ceaCollectExceptionAnnotationMechanism = fmap SomeExceptionAnnotation collector
+        }
+  _ <- atomicModifyIORef'_ collectExceptionAnnotationMechanismRef (const cea)
+  return ()
+
 -- | A collection of backtraces.
 data Backtraces =
     Backtraces {
@@ -124,6 +155,14 @@ displayBacktraces bts = concat
 instance ExceptionAnnotation Backtraces where
     displayExceptionAnnotation = displayBacktraces
 
+-- | Collect 'SomeExceptionAnnotation' based on the configuration of the
+-- global 'CollectExceptionAnnotationMechanism'.
+--
+collectExceptionAnnotation :: HasCallStack => IO SomeExceptionAnnotation
+collectExceptionAnnotation = HCS.withFrozenCallStack $ do
+  cea <- getCollectExceptionAnnotationMechanism
+  ceaCollectExceptionAnnotationMechanism cea
+
 -- | Collect a set of 'Backtraces'.
 collectBacktraces :: (?callStack :: CallStack) => IO Backtraces
 collectBacktraces = HCS.withFrozenCallStack $ do

libraries/ghc-internal/src/GHC/Internal/Exception/Backtrace.hs-boot

@@ -5,11 +5,7 @@ module GHC.Internal.Exception.Backtrace where
 
 import GHC.Internal.Base (IO)
 import GHC.Internal.Stack.Types (HasCallStack)
-import GHC.Internal.Exception.Context (ExceptionAnnotation)
-
-data Backtraces
-
-instance ExceptionAnnotation Backtraces
+import GHC.Internal.Exception.Context (SomeExceptionAnnotation)
 
 -- For GHC.Exception
-collectBacktraces :: HasCallStack => IO Backtraces
+collectExceptionAnnotation :: HasCallStack => IO SomeExceptionAnnotation

libraries/ghc-internal/tests/Makefile

@@ -0,0 +1,7 @@
+# This Makefile runs the tests using GHC's testsuite framework.  It
+# assumes the package is part of a GHC build tree with the testsuite
+# installed in ../../../testsuite.
+
+TOP=../../../testsuite
+include $(TOP)/mk/boilerplate.mk
+include $(TOP)/mk/test.mk

libraries/ghc-internal/tests/all.T

@@ -0,0 +1,7 @@
+# This Makefile runs the tests using GHC's testsuite framework.  It
+# assumes the package is part of a GHC build tree with the testsuite
+# installed in ../../../testsuite.
+
+TOP=../../../../testsuite
+include $(TOP)/mk/boilerplate.mk
+include $(TOP)/mk/test.mk

libraries/ghc-internal/tests/backtraces/Makefile

@@ -0,0 +1,27 @@
+import Control.Exception
+import Control.Exception.Annotation
+import GHC.Internal.Exception.Backtrace
+
+data MyBacktraces = MyBacktraces
+
+instance ExceptionAnnotation MyBacktraces where
+  displayExceptionAnnotation MyBacktraces = "MyBacktraces"
+
+main :: IO ()
+main = do
+  setCollectExceptionAnnotation (pure MyBacktraces)
+  catchAndPrint functionThatThrows
+  catchAndPrint functionThatErrors
+
+catchAndPrint :: IO () -> IO ()
+catchAndPrint act = do
+  act `catch` \(exc :: SomeException) -> do
+    putStrLn $ displayExceptionWithInfo exc
+
+functionThatThrows :: IO ()
+functionThatThrows = do
+  throwIO $ ErrorCall "functionThatThrows"
+
+functionThatErrors :: a
+functionThatErrors = do
+  error "functionThatErrors"

libraries/ghc-internal/tests/backtraces/T14532a.hs

@@ -0,0 +1,10 @@
+ghc-internal:GHC.Internal.Exception.ErrorCall:
+
+functionThatThrows
+
+MyBacktraces
+ghc-internal:GHC.Internal.Exception.ErrorCall:
+
+functionThatErrors
+
+MyBacktraces

libraries/ghc-internal/tests/backtraces/T14532a.stdout

@@ -0,0 +1,41 @@
+import Control.Exception
+import Control.Exception.Annotation
+import GHC.Internal.Exception.Backtrace
+
+
+newtype MyBacktraces = MyBacktraces Backtraces
+
+instance ExceptionAnnotation MyBacktraces where
+  displayExceptionAnnotation (MyBacktraces bt) = unlines
+    [ "My custom Backtraces:"
+    , displayBacktraces bt
+    ]
+
+main :: IO ()
+main = do
+  setCollectExceptionAnnotation (MyBacktraces <$> collectBacktraces)
+  catchAndPrint whileHandlingThrowIO
+  catchAndPrint whileHandlingError
+
+catchAndPrint :: IO () -> IO ()
+catchAndPrint act = do
+  act `catch` \(exc :: SomeException) -> do
+    putStrLn $ displayExceptionWithInfo exc
+
+whileHandlingThrowIO :: IO ()
+whileHandlingThrowIO =
+  handleJust
+    (\e -> case e of
+      ErrorCall{} -> Just e
+      _ -> Nothing)
+    (\_ -> throwIO $ ErrorCall "Error in Exception Handler")
+    (throwIO $ ErrorCall "Main Error")
+
+whileHandlingError :: IO ()
+whileHandlingError = do
+  handleJust
+    (\e -> case e of
+      ErrorCall{} -> Just e
+      _ -> Nothing)
+    (\_ -> error "Error in Exception Handler")
+    (error "Main Error")