-- SPDX-FileCopyrightText: 2022 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

{-# OPTIONS_HADDOCK not-home #-}

-- | This module defines annotated exceptions and machinery to work with them.
module Test.Cleveland.Internal.Exceptions.Annotated
  ( module Test.Cleveland.Internal.Exceptions.Annotated
  ) where

import Data.Dependent.Map qualified as Map
import Data.GADT.Compare (GCompare(..), GEq(..), GOrdering(..))
import Data.Typeable (cast, eqT)
import Fmt (Buildable(..), Builder, build, pretty)
import Text.Show qualified as Show (Show(..))
import Type.Reflection (typeRep)

-- $setup
-- >>> import Test.Cleveland.Internal.Abstract
-- >>> import Test.Cleveland.Internal.Exceptions.CallStack

----------------------------------------------------------------------------
-- AnnotatedException
----------------------------------------------------------------------------

-- | 'SomeException' augmented with arbitrary annotations.
data AnnotatedException = AnnotatedException AnnotationMap SomeException
  deriving stock Int -> AnnotatedException -> ShowS
[AnnotatedException] -> ShowS
AnnotatedException -> String
(Int -> AnnotatedException -> ShowS)
-> (AnnotatedException -> String)
-> ([AnnotatedException] -> ShowS)
-> Show AnnotatedException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [AnnotatedException] -> ShowS
$cshowList :: [AnnotatedException] -> ShowS
show :: AnnotatedException -> String
$cshow :: AnnotatedException -> String
showsPrec :: Int -> AnnotatedException -> ShowS
$cshowsPrec :: Int -> AnnotatedException -> ShowS
Show

instance Exception AnnotatedException where
  displayException :: AnnotatedException -> String
displayException (AnnotatedException AnnotationMap
annmap SomeException
e) =
    Builder -> String
forall a b. (Buildable a, FromBuilder b) => a -> b
pretty (Builder -> String) -> Builder -> String
forall a b. (a -> b) -> a -> b
$ AnnotationMap -> String -> Builder
forall a. Buildable a => AnnotationMap -> a -> Builder
printAnnotations AnnotationMap
annmap (SomeException -> String
forall e. Exception e => e -> String
displayException SomeException
e)

-- | Insert a new annotation into 'AnnotatedException'.
insertAnnEx'
  :: (Semigroup ann, ExceptionAnnotation ann)
  => ann -> AnnotatedException -> AnnotatedException
insertAnnEx' :: forall ann.
(Semigroup ann, ExceptionAnnotation ann) =>
ann -> AnnotatedException -> AnnotatedException
insertAnnEx' ann
ann (AnnotatedException AnnotationMap
anns SomeException
ex) = AnnotationMap -> SomeException -> AnnotatedException
AnnotatedException (ann -> AnnotationMap -> AnnotationMap
forall t.
(Semigroup t, ExceptionAnnotation t) =>
t -> AnnotationMap -> AnnotationMap
insertAnn ann
ann AnnotationMap
anns) SomeException
ex

-- | Insert a new annotation into 'SomeException'.
--
-- If it's already an 'AnnotatedException', simply adds a new annotation.
-- Otherwise, wraps the exception into 'AnnotatedException' first.
insertAnnEx
  :: (Semigroup ann, ExceptionAnnotation ann)
  => ann -> SomeException -> SomeException
insertAnnEx :: forall ann.
(Semigroup ann, ExceptionAnnotation ann) =>
ann -> SomeException -> SomeException
insertAnnEx ann
ann SomeException
se = AnnotatedException -> SomeException
forall e. Exception e => e -> SomeException
toException (AnnotatedException -> SomeException)
-> AnnotatedException -> SomeException
forall a b. (a -> b) -> a -> b
$! case forall e. Exception e => SomeException -> Maybe e
fromException @AnnotatedException SomeException
se of
  Just AnnotatedException
ex -> ann -> AnnotatedException -> AnnotatedException
forall ann.
(Semigroup ann, ExceptionAnnotation ann) =>
ann -> AnnotatedException -> AnnotatedException
insertAnnEx' ann
ann AnnotatedException
ex
  Maybe AnnotatedException
Nothing -> AnnotationMap -> SomeException -> AnnotatedException
AnnotatedException (ann -> AnnotationMap
forall t. ExceptionAnnotation t => t -> AnnotationMap
singletonAnn ann
ann) SomeException
se

-- | Remove an annotation of the given type from 'AnnotatedException'.
removeAnnEx' :: forall ann. ExceptionAnnotation ann => AnnotatedException -> AnnotatedException
removeAnnEx' :: forall ann.
ExceptionAnnotation ann =>
AnnotatedException -> AnnotatedException
removeAnnEx' (AnnotatedException AnnotationMap
anns SomeException
ex) = AnnotationMap -> SomeException -> AnnotatedException
AnnotatedException (forall t. ExceptionAnnotation t => AnnotationMap -> AnnotationMap
removeAnn @ann AnnotationMap
anns) SomeException
ex

-- | Remove an annotation of the given type from 'SomeException'. If it's not an
-- 'AnnotatedException', this is a no-op.
removeAnnEx :: forall ann. ExceptionAnnotation ann => SomeException -> SomeException
removeAnnEx :: forall ann.
ExceptionAnnotation ann =>
SomeException -> SomeException
removeAnnEx SomeException
se = case forall e. Exception e => SomeException -> Maybe e
fromException @AnnotatedException SomeException
se of
  Just AnnotatedException
ex -> AnnotatedException -> SomeException
forall e. Exception e => e -> SomeException
toException (AnnotatedException -> SomeException)
-> AnnotatedException -> SomeException
forall a b. (a -> b) -> a -> b
$ forall ann.
ExceptionAnnotation ann =>
AnnotatedException -> AnnotatedException
removeAnnEx' @ann AnnotatedException
ex
  Maybe AnnotatedException
Nothing -> SomeException
se

-- | Extract the given annotation type from 'AnnotatedException' if possible.
lookupAnnEx' :: forall ann. (ExceptionAnnotation ann) => AnnotatedException -> Maybe ann
lookupAnnEx' :: forall ann.
ExceptionAnnotation ann =>
AnnotatedException -> Maybe ann
lookupAnnEx' (AnnotatedException AnnotationMap
anns SomeException
_) = forall t. ExceptionAnnotation t => AnnotationMap -> Maybe t
lookupAnn @ann AnnotationMap
anns

-- | Extract the given annotation type from 'SomeException' if possible.
lookupAnnEx :: forall ann. (ExceptionAnnotation ann) => SomeException -> Maybe ann
lookupAnnEx :: forall ann. ExceptionAnnotation ann => SomeException -> Maybe ann
lookupAnnEx SomeException
se = AnnotatedException -> Maybe ann
forall ann.
ExceptionAnnotation ann =>
AnnotatedException -> Maybe ann
lookupAnnEx' (AnnotatedException -> Maybe ann)
-> Maybe AnnotatedException -> Maybe ann
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< forall e. Exception e => SomeException -> Maybe e
fromException @AnnotatedException SomeException
se

{- | Try to extract a given exception type, skipping over annotations. This
essentially makes 'AnnotatedException' transparent for this function.

This is intended to be used as the default implementation of 'fromException' for
exception types that are expected to be annotated, as to make
'AnnotatedException' transparent for 'catch', 'try', etc.

For example:

@
instance Exception GenericTestError where
  displayException = pretty
  fromException = fromPossiblyAnnotatedException
@

Now, when catching @GenericTestError@, it'll work even if wrapped in
'AnnotatedException':

>>> try @_ @GenericTestError (addCallStack $ throwM UnexpectedSuccess)
Left UnexpectedSuccess

When catching exceptions not explicitly set up like this, this trick
unfortunately won't work, so using this function explicitly is required:

>>> data ExceptionToCatch = ExceptionToCatch deriving stock Show
>>> instance Exception ExceptionToCatch
>>> action = addCallStack $ throwM ExceptionToCatch
>>> try @_ @ExceptionToCatch action
*** Exception: AnnotatedException ...
...ExceptionToCatch
>>> :{
action `catch` \(err :: SomeException) -> case fromPossiblyAnnotatedException err of
  Just (err' :: ExceptionToCatch) -> pass
  Nothing -> throwM err
:}
-}
fromPossiblyAnnotatedException :: Exception e => SomeException -> Maybe e
fromPossiblyAnnotatedException :: forall e. Exception e => SomeException -> Maybe e
fromPossiblyAnnotatedException SomeException
se =
  (do
    AnnotatedException AnnotationMap
_ SomeException
e <- SomeException -> Maybe AnnotatedException
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
se
    SomeException -> Maybe e
forall e. Exception e => SomeException -> Maybe e
fromException SomeException
e
  ) Maybe e -> Maybe e -> Maybe e
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> SomeException -> Maybe e
forall {b}. Typeable b => SomeException -> Maybe b
fromExceptionDefault SomeException
se
  where
    fromExceptionDefault :: SomeException -> Maybe b
fromExceptionDefault (SomeException e
e) = e -> Maybe b
forall a b. (Typeable a, Typeable b) => a -> Maybe b
cast e
e

----------------------------------------------------------------------------
-- AnnotationMap
----------------------------------------------------------------------------

data Key ann where
  Key :: ExceptionAnnotation a => Key a

instance GEq Key where
  geq :: forall a b. Key a -> Key b -> Maybe (a :~: b)
geq (Key a
Key :: Key a) (Key b
Key :: Key b) = forall {k} (a :: k) (b :: k).
(Typeable a, Typeable b) =>
Maybe (a :~: b)
forall a b. (Typeable a, Typeable b) => Maybe (a :~: b)
eqT @a @b

instance GCompare Key where
  gcompare :: forall a b. Key a -> Key b -> GOrdering a b
gcompare (Key a
Key :: Key a) (Key b
Key :: Key b) =
    case Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (forall ann. ExceptionAnnotation ann => Int
annotationPriority @a) (forall ann. ExceptionAnnotation ann => Int
annotationPriority @b) of
      Ordering
LT -> GOrdering a b
forall {k} (a :: k) (b :: k). GOrdering a b
GLT
      Ordering
GT -> GOrdering a b
forall {k} (a :: k) (b :: k). GOrdering a b
GGT
      Ordering
EQ -> TypeRep a -> TypeRep b -> GOrdering a b
forall k (f :: k -> *) (a :: k) (b :: k).
GCompare f =>
f a -> f b -> GOrdering a b
gcompare (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @a) (forall a. Typeable a => TypeRep a
forall {k} (a :: k). Typeable a => TypeRep a
typeRep @b)

-- | A map from 'Key' to corresponding annotation.
newtype AnnotationMap = AnnotationMap { AnnotationMap -> DMap Key Identity
unAnnotationMap :: Map.DMap Key Identity }

instance Show AnnotationMap where
  show :: AnnotationMap -> String
show = (forall v. Key v -> Identity v -> String)
-> Bool -> Bool -> DMap Key Identity -> String
forall {k1} (k2 :: k1 -> *) (f :: k1 -> *).
(forall (v :: k1). k2 v -> f v -> String)
-> Bool -> Bool -> DMap k2 f -> String
Map.showTreeWith (\Key v
Key -> Identity v -> String
forall a. Show a => a -> String
Show.show) Bool
True Bool
False (DMap Key Identity -> String)
-> (AnnotationMap -> DMap Key Identity) -> AnnotationMap -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. AnnotationMap -> DMap Key Identity
unAnnotationMap

-- | Print all annotations from an 'AnnotationMap', given the initial message.
--
-- Essentially right-fold using 'displayAnnotation'.
printAnnotations :: Buildable a => AnnotationMap -> a -> Builder
printAnnotations :: forall a. Buildable a => AnnotationMap -> a -> Builder
printAnnotations (AnnotationMap DMap Key Identity
annmap) = (Builder -> DMap Key Identity -> Builder)
-> DMap Key Identity -> Builder -> Builder
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((forall v. Key v -> Identity v -> Builder -> Builder)
-> Builder -> DMap Key Identity -> Builder
forall {k1} (k2 :: k1 -> *) (f :: k1 -> *) b.
(forall (v :: k1). k2 v -> f v -> b -> b) -> b -> DMap k2 f -> b
Map.foldrWithKey forall v. Key v -> Identity v -> Builder -> Builder
go) DMap Key Identity
annmap (Builder -> Builder) -> (a -> Builder) -> a -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Builder
forall p. Buildable p => p -> Builder
build
  where
    go :: Key t -> Identity t -> Builder -> Builder
    go :: forall v. Key v -> Identity v -> Builder -> Builder
go Key t
Key (Identity t
a) Builder
acc = t -> Builder -> Builder
forall ann. ExceptionAnnotation ann => ann -> Builder -> Builder
displayAnnotation t
a Builder
acc

-- | Try to find the given type @t@ in an 'AnnotationMap'.
lookupAnn :: ExceptionAnnotation t => AnnotationMap -> Maybe t
lookupAnn :: forall t. ExceptionAnnotation t => AnnotationMap -> Maybe t
lookupAnn (AnnotationMap DMap Key Identity
anns) = Identity t -> t
forall a. Identity a -> a
runIdentity (Identity t -> t) -> Maybe (Identity t) -> Maybe t
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Key t -> DMap Key Identity -> Maybe (Identity t)
forall {k1} (k2 :: k1 -> *) (f :: k1 -> *) (v :: k1).
GCompare k2 =>
k2 v -> DMap k2 f -> Maybe (f v)
Map.lookup Key t
forall a. ExceptionAnnotation a => Key a
Key DMap Key Identity
anns

-- | Construct a singleton 'AnnotationMap'.
singletonAnn :: ExceptionAnnotation t => t -> AnnotationMap
singletonAnn :: forall t. ExceptionAnnotation t => t -> AnnotationMap
singletonAnn = DMap Key Identity -> AnnotationMap
AnnotationMap (DMap Key Identity -> AnnotationMap)
-> (t -> DMap Key Identity) -> t -> AnnotationMap
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Key t -> Identity t -> DMap Key Identity
forall {k1} (k2 :: k1 -> *) (v :: k1) (f :: k1 -> *).
k2 v -> f v -> DMap k2 f
Map.singleton Key t
forall a. ExceptionAnnotation a => Key a
Key (Identity t -> DMap Key Identity)
-> (t -> Identity t) -> t -> DMap Key Identity
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Identity t
forall a. a -> Identity a
Identity

-- | Insert a new annotation into an 'AnnotationMap'. If an annotation of this
-- type already exists, merge the annotations using the 'Semigroup' instance.
insertAnn :: (Semigroup t, ExceptionAnnotation t) => t -> AnnotationMap -> AnnotationMap
insertAnn :: forall t.
(Semigroup t, ExceptionAnnotation t) =>
t -> AnnotationMap -> AnnotationMap
insertAnn t
x (AnnotationMap DMap Key Identity
anns) = DMap Key Identity -> AnnotationMap
AnnotationMap (DMap Key Identity -> AnnotationMap)
-> DMap Key Identity -> AnnotationMap
forall a b. (a -> b) -> a -> b
$ (Identity t -> Identity t -> Identity t)
-> Key t -> Identity t -> DMap Key Identity -> DMap Key Identity
forall {k1} (k2 :: k1 -> *) (f :: k1 -> *) (v :: k1).
GCompare k2 =>
(f v -> f v -> f v) -> k2 v -> f v -> DMap k2 f -> DMap k2 f
Map.insertWith Identity t -> Identity t -> Identity t
forall a. Semigroup a => a -> a -> a
(<>) Key t
forall a. ExceptionAnnotation a => Key a
Key (t -> Identity t
forall a. a -> Identity a
Identity t
x) DMap Key Identity
anns

-- | Remove an annotation of a given type @t@ from 'AnnotationMap'.
removeAnn :: forall t. ExceptionAnnotation t => AnnotationMap -> AnnotationMap
removeAnn :: forall t. ExceptionAnnotation t => AnnotationMap -> AnnotationMap
removeAnn (AnnotationMap DMap Key Identity
am) = DMap Key Identity -> AnnotationMap
AnnotationMap (DMap Key Identity -> AnnotationMap)
-> DMap Key Identity -> AnnotationMap
forall a b. (a -> b) -> a -> b
$ Key t -> DMap Key Identity -> DMap Key Identity
forall {k1} (k2 :: k1 -> *) (f :: k1 -> *) (v :: k1).
GCompare k2 =>
k2 v -> DMap k2 f -> DMap k2 f
Map.delete (forall a. ExceptionAnnotation a => Key a
Key @t) DMap Key Identity
am

----------------------------------------------------------------------------
-- ExceptionAnnotation
----------------------------------------------------------------------------

-- | Type class for exception annotations.
class (Show ann, Typeable ann) => ExceptionAnnotation ann where
  -- | Given an annotation and the error message (as 'Builder'), produce the
  -- text annotated with the annotation.
  displayAnnotation :: ann -> Builder -> Builder

  -- | Relative priority for sorting annotations. Annotations with higher
  -- priority will be applied first. Default is @0@.
  annotationPriority :: Int
  annotationPriority = Int
0

-- | Add an annotation to exceptions thrown by a monadic action.
annotateExceptions
  :: forall m a ann. (MonadCatch m, Semigroup ann, ExceptionAnnotation ann)
  => ann -> m a -> m a
annotateExceptions :: forall (m :: * -> *) a ann.
(MonadCatch m, Semigroup ann, ExceptionAnnotation ann) =>
ann -> m a -> m a
annotateExceptions ann
ann m a
ma = m a
ma m a -> (SomeException -> m a) -> m a
forall (m :: * -> *) e a.
(MonadCatch m, Exception e) =>
m a -> (e -> m a) -> m a
`catch` (SomeException -> m a
forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (SomeException -> m a)
-> (SomeException -> SomeException) -> SomeException -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ann -> SomeException -> SomeException
forall ann.
(Semigroup ann, ExceptionAnnotation ann) =>
ann -> SomeException -> SomeException
insertAnnEx ann
ann)