    -- |
    -- Support for source code annotation feature of GHC. That is the ANN pragma.
    --
    -- (c) The University of Glasgow 2006
    -- (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
    --
    {-# LANGUAGE DeriveFunctor #-}
    module GHC.Types.Annotations (
            -- * Main Annotation data types
            Annotation(..), AnnPayload,
            AnnTarget(..), CoreAnnTarget,

            -- * AnnEnv for collecting and querying Annotations
            AnnEnv,
            mkAnnEnv, extendAnnEnvList, plusAnnEnv, emptyAnnEnv,
            findAnns, findAnnsByTypeRep,
            deserializeAnns
        ) where

    import GHC.Prelude

    import GHC.Utils.Binary
    import GHC.Unit.Module ( Module )
    import GHC.Unit.Module.Env
    import GHC.Types.Name.Env
    import GHC.Types.Name
    import GHC.Utils.Outputable
    import GHC.Serialized

    import Control.Monad
    import Data.Maybe
    import Data.Typeable
    import Data.Word        ( Word8 )

    -- | Represents an annotation after it has been sufficiently desugared from
    -- it's initial form of 'GHC.Hs.Decls.AnnDecl'
    data Annotation = Annotation {
            Annotation -> CoreAnnTarget
    ann_target :: CoreAnnTarget,    -- ^ The target of the annotation
            Annotation -> AnnPayload
    ann_value  :: AnnPayload
        }

    type AnnPayload = Serialized    -- ^ The "payload" of an annotation
                                    --   allows recovery of its value at a given type,
                                    --   and can be persisted to an interface file

    -- | An annotation target
    data AnnTarget name
      = NamedTarget name          -- ^ We are annotating something with a name:
                                  --      a type or identifier
      | ModuleTarget Module       -- ^ We are annotating a particular module
      deriving ((forall a b. (a -> b) -> AnnTarget a -> AnnTarget b)
    -> (forall a b. a -> AnnTarget b -> AnnTarget a)
    -> Functor AnnTarget
    forall a b. a -> AnnTarget b -> AnnTarget a
    forall a b. (a -> b) -> AnnTarget a -> AnnTarget b
    forall (f :: * -> *).
    (forall a b. (a -> b) -> f a -> f b)
    -> (forall a b. a -> f b -> f a) -> Functor f
    <$ :: forall a b. a -> AnnTarget b -> AnnTarget a
    $c<$ :: forall a b. a -> AnnTarget b -> AnnTarget a
    fmap :: forall a b. (a -> b) -> AnnTarget a -> AnnTarget b
    $cfmap :: forall a b. (a -> b) -> AnnTarget a -> AnnTarget b
    Functor)

    -- | The kind of annotation target found in the middle end of the compiler
    type CoreAnnTarget = AnnTarget Name

    instance Outputable name => Outputable (AnnTarget name) where
        ppr :: AnnTarget name -> SDoc
    ppr (NamedTarget name
    nm) = String -> SDoc
    text String
    "Named target" SDoc -> SDoc -> SDoc
    <+> name -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr name
    nm
        ppr (ModuleTarget Module
    mod) = String -> SDoc
    text String
    "Module target" SDoc -> SDoc -> SDoc
    <+> Module -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr Module
    mod

    instance Binary name => Binary (AnnTarget name) where
        put_ :: BinHandle -> AnnTarget name -> IO ()
    put_ BinHandle
    bh (NamedTarget name
    a) = do
            BinHandle -> Word8 -> IO ()
    putByte BinHandle
    bh Word8
    0
            BinHandle -> name -> IO ()
    forall a. Binary a => BinHandle -> a -> IO ()
    put_ BinHandle
    bh name
    a
        put_ BinHandle
    bh (ModuleTarget Module
    a) = do
            BinHandle -> Word8 -> IO ()
    putByte BinHandle
    bh Word8
    1
            BinHandle -> Module -> IO ()
    forall a. Binary a => BinHandle -> a -> IO ()
    put_ BinHandle
    bh Module
    a
        get :: BinHandle -> IO (AnnTarget name)
    get BinHandle
    bh = do
            Word8
    h <- BinHandle -> IO Word8
    getByte BinHandle
    bh
            case Word8
    h of
                Word8
    0 -> (name -> AnnTarget name) -> IO name -> IO (AnnTarget name)
    forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
    liftM name -> AnnTarget name
    forall name. name -> AnnTarget name
    NamedTarget  (IO name -> IO (AnnTarget name)) -> IO name -> IO (AnnTarget name)
    forall a b. (a -> b) -> a -> b
    $ BinHandle -> IO name
    forall a. Binary a => BinHandle -> IO a
    get BinHandle
    bh
                Word8
    _ -> (Module -> AnnTarget name) -> IO Module -> IO (AnnTarget name)
    forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
    liftM Module -> AnnTarget name
    forall name. Module -> AnnTarget name
    ModuleTarget (IO Module -> IO (AnnTarget name))
    -> IO Module -> IO (AnnTarget name)
    forall a b. (a -> b) -> a -> b
    $ BinHandle -> IO Module
    forall a. Binary a => BinHandle -> IO a
    get BinHandle
    bh

    instance Outputable Annotation where
        ppr :: Annotation -> SDoc
    ppr Annotation
    ann = CoreAnnTarget -> SDoc
    forall a. Outputable a => a -> SDoc
    ppr (Annotation -> CoreAnnTarget
    ann_target Annotation
    ann)

    -- | A collection of annotations
    data AnnEnv = MkAnnEnv { AnnEnv -> ModuleEnv [AnnPayload]
    ann_mod_env :: !(ModuleEnv [AnnPayload])
                           , AnnEnv -> NameEnv [AnnPayload]
    ann_name_env :: !(NameEnv [AnnPayload])
                           }

    -- | An empty annotation environment.
    emptyAnnEnv :: AnnEnv
    emptyAnnEnv :: AnnEnv
    emptyAnnEnv = ModuleEnv [AnnPayload] -> NameEnv [AnnPayload] -> AnnEnv
    MkAnnEnv ModuleEnv [AnnPayload]
    forall a. ModuleEnv a
    emptyModuleEnv NameEnv [AnnPayload]
    forall a. NameEnv a
    emptyNameEnv

    -- | Construct a new annotation environment that contains the list of
    -- annotations provided.
    mkAnnEnv :: [Annotation] -> AnnEnv
    mkAnnEnv :: [Annotation] -> AnnEnv
    mkAnnEnv = AnnEnv -> [Annotation] -> AnnEnv
    extendAnnEnvList AnnEnv
    emptyAnnEnv

    -- | Add the given annotation to the environment.
    extendAnnEnvList :: AnnEnv -> [Annotation] -> AnnEnv
    extendAnnEnvList :: AnnEnv -> [Annotation] -> AnnEnv
    extendAnnEnvList AnnEnv
    env =
      (AnnEnv -> Annotation -> AnnEnv)
    -> AnnEnv -> [Annotation] -> AnnEnv
    forall (t :: * -> *) b a.
    Foldable t =>
    (b -> a -> b) -> b -> t a -> b
    foldl' AnnEnv -> Annotation -> AnnEnv
    extendAnnEnv AnnEnv
    env

    extendAnnEnv :: AnnEnv -> Annotation -> AnnEnv
    extendAnnEnv :: AnnEnv -> Annotation -> AnnEnv
    extendAnnEnv (MkAnnEnv ModuleEnv [AnnPayload]
    mod_env NameEnv [AnnPayload]
    name_env) (Annotation CoreAnnTarget
    tgt AnnPayload
    payload) =
      case CoreAnnTarget
    tgt of
        NamedTarget Name
    name -> ModuleEnv [AnnPayload] -> NameEnv [AnnPayload] -> AnnEnv
    MkAnnEnv ModuleEnv [AnnPayload]
    mod_env (([AnnPayload] -> [AnnPayload] -> [AnnPayload])
    -> NameEnv [AnnPayload]
    -> Name
    -> [AnnPayload]
    -> NameEnv [AnnPayload]
    forall a. (a -> a -> a) -> NameEnv a -> Name -> a -> NameEnv a
    extendNameEnv_C [AnnPayload] -> [AnnPayload] -> [AnnPayload]
    forall a. [a] -> [a] -> [a]
    (++) NameEnv [AnnPayload]
    name_env Name
    name [AnnPayload
    payload])
        ModuleTarget Module
    mod -> ModuleEnv [AnnPayload] -> NameEnv [AnnPayload] -> AnnEnv
    MkAnnEnv (([AnnPayload] -> [AnnPayload] -> [AnnPayload])
    -> ModuleEnv [AnnPayload]
    -> Module
    -> [AnnPayload]
    -> ModuleEnv [AnnPayload]
    forall a.
    (a -> a -> a) -> ModuleEnv a -> Module -> a -> ModuleEnv a
    extendModuleEnvWith [AnnPayload] -> [AnnPayload] -> [AnnPayload]
    forall a. [a] -> [a] -> [a]
    (++) ModuleEnv [AnnPayload]
    mod_env Module
    mod [AnnPayload
    payload]) NameEnv [AnnPayload]
    name_env

    -- | Union two annotation environments.
    plusAnnEnv :: AnnEnv -> AnnEnv -> AnnEnv
    plusAnnEnv :: AnnEnv -> AnnEnv -> AnnEnv
    plusAnnEnv AnnEnv
    a AnnEnv
    b =
      MkAnnEnv { ann_mod_env :: ModuleEnv [AnnPayload]
    ann_mod_env = ([AnnPayload] -> [AnnPayload] -> [AnnPayload])
    -> ModuleEnv [AnnPayload]
    -> ModuleEnv [AnnPayload]
    -> ModuleEnv [AnnPayload]
    forall a.
    (a -> a -> a) -> ModuleEnv a -> ModuleEnv a -> ModuleEnv a
    plusModuleEnv_C [AnnPayload] -> [AnnPayload] -> [AnnPayload]
    forall a. [a] -> [a] -> [a]
    (++) (AnnEnv -> ModuleEnv [AnnPayload]
    ann_mod_env AnnEnv
    a) (AnnEnv -> ModuleEnv [AnnPayload]
    ann_mod_env AnnEnv
    b)
               , ann_name_env :: NameEnv [AnnPayload]
    ann_name_env = ([AnnPayload] -> [AnnPayload] -> [AnnPayload])
    -> NameEnv [AnnPayload]
    -> NameEnv [AnnPayload]
    -> NameEnv [AnnPayload]
    forall a. (a -> a -> a) -> NameEnv a -> NameEnv a -> NameEnv a
    plusNameEnv_C [AnnPayload] -> [AnnPayload] -> [AnnPayload]
    forall a. [a] -> [a] -> [a]
    (++) (AnnEnv -> NameEnv [AnnPayload]
    ann_name_env AnnEnv
    a) (AnnEnv -> NameEnv [AnnPayload]
    ann_name_env AnnEnv
    b)
               }

    -- | Find the annotations attached to the given target as 'Typeable'
    --   values of your choice. If no deserializer is specified,
    --   only transient annotations will be returned.
    findAnns :: Typeable a => ([Word8] -> a) -> AnnEnv -> CoreAnnTarget -> [a]
    findAnns :: forall a.
    Typeable a =>
    ([Word8] -> a) -> AnnEnv -> CoreAnnTarget -> [a]
    findAnns [Word8] -> a
    deserialize AnnEnv
    env
      = (AnnPayload -> Maybe a) -> [AnnPayload] -> [a]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe (([Word8] -> a) -> AnnPayload -> Maybe a
    forall a. Typeable a => ([Word8] -> a) -> AnnPayload -> Maybe a
    fromSerialized [Word8] -> a
    deserialize) ([AnnPayload] -> [a])
    -> (CoreAnnTarget -> [AnnPayload]) -> CoreAnnTarget -> [a]
    forall b c a. (b -> c) -> (a -> b) -> a -> c
    . AnnEnv -> CoreAnnTarget -> [AnnPayload]
    findAnnPayloads AnnEnv
    env

    -- | Find the annotations attached to the given target as 'Typeable'
    --   values of your choice. If no deserializer is specified,
    --   only transient annotations will be returned.
    findAnnsByTypeRep :: AnnEnv -> CoreAnnTarget -> TypeRep -> [[Word8]]
    findAnnsByTypeRep :: AnnEnv -> CoreAnnTarget -> TypeRep -> [[Word8]]
    findAnnsByTypeRep AnnEnv
    env CoreAnnTarget
    target TypeRep
    tyrep
      = [ [Word8]
    ws | Serialized TypeRep
    tyrep' [Word8]
    ws <- AnnEnv -> CoreAnnTarget -> [AnnPayload]
    findAnnPayloads AnnEnv
    env CoreAnnTarget
    target
        , TypeRep
    tyrep' TypeRep -> TypeRep -> Bool
    forall a. Eq a => a -> a -> Bool
    == TypeRep
    tyrep ]

    -- | Find payloads for the given 'CoreAnnTarget' in an 'AnnEnv'.
    findAnnPayloads :: AnnEnv -> CoreAnnTarget -> [AnnPayload]
    findAnnPayloads :: AnnEnv -> CoreAnnTarget -> [AnnPayload]
    findAnnPayloads AnnEnv
    env CoreAnnTarget
    target =
      case CoreAnnTarget
    target of
        ModuleTarget Module
    mod -> ModuleEnv [AnnPayload] -> [AnnPayload] -> Module -> [AnnPayload]
    forall a. ModuleEnv a -> a -> Module -> a
    lookupWithDefaultModuleEnv (AnnEnv -> ModuleEnv [AnnPayload]
    ann_mod_env AnnEnv
    env) [] Module
    mod
        NamedTarget Name
    name -> [AnnPayload] -> Maybe [AnnPayload] -> [AnnPayload]
    forall a. a -> Maybe a -> a
    fromMaybe [] (Maybe [AnnPayload] -> [AnnPayload])
    -> Maybe [AnnPayload] -> [AnnPayload]
    forall a b. (a -> b) -> a -> b
    $ NameEnv [AnnPayload] -> Name -> Maybe [AnnPayload]
    forall a. NameEnv a -> Name -> Maybe a
    lookupNameEnv (AnnEnv -> NameEnv [AnnPayload]
    ann_name_env AnnEnv
    env) Name
    name

    -- | Deserialize all annotations of a given type. This happens lazily, that is
    --   no deserialization will take place until the [a] is actually demanded and
    --   the [a] can also be empty (the UniqFM is not filtered).
    deserializeAnns :: Typeable a => ([Word8] -> a) -> AnnEnv -> (ModuleEnv [a], NameEnv [a])
    deserializeAnns :: forall a.
    Typeable a =>
    ([Word8] -> a) -> AnnEnv -> (ModuleEnv [a], NameEnv [a])
    deserializeAnns [Word8] -> a
    deserialize AnnEnv
    env
      = ( ([AnnPayload] -> [a]) -> ModuleEnv [AnnPayload] -> ModuleEnv [a]
    forall a b. (a -> b) -> ModuleEnv a -> ModuleEnv b
    mapModuleEnv [AnnPayload] -> [a]
    deserAnns (AnnEnv -> ModuleEnv [AnnPayload]
    ann_mod_env AnnEnv
    env)
        , ([AnnPayload] -> [a]) -> NameEnv [AnnPayload] -> NameEnv [a]
    forall elt1 elt2. (elt1 -> elt2) -> NameEnv elt1 -> NameEnv elt2
    mapNameEnv [AnnPayload] -> [a]
    deserAnns (AnnEnv -> NameEnv [AnnPayload]
    ann_name_env AnnEnv
    env)
        )
      where deserAnns :: [AnnPayload] -> [a]
    deserAnns = (AnnPayload -> Maybe a) -> [AnnPayload] -> [a]
    forall a b. (a -> Maybe b) -> [a] -> [b]
    mapMaybe (([Word8] -> a) -> AnnPayload -> Maybe a
    forall a. Typeable a => ([Word8] -> a) -> AnnPayload -> Maybe a
    fromSerialized [Word8] -> a
    deserialize)