-- |
-- 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 :: ModuleEnv [AnnPayload] -> NameEnv [AnnPayload] -> AnnEnv
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)