GHC/Types/Annotations.hs

    -- |
    -- 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 {
            ann_target :: CoreAnnTarget,    -- ^ The target of the annotation
            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 (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 (NamedTarget nm) = text "Named target" <+> ppr nm
        ppr (ModuleTarget mod) = text "Module target" <+> ppr mod

    instance Binary name => Binary (AnnTarget name) where
        put_ bh (NamedTarget a) = do
            putByte bh 0
            put_ bh a
        put_ bh (ModuleTarget a) = do
            putByte bh 1
            put_ bh a
        get bh = do
            h <- getByte bh
            case h of
                0 -> liftM NamedTarget  $ get bh
                _ -> liftM ModuleTarget $ get bh

    instance Outputable Annotation where
        ppr ann = ppr (ann_target ann)

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

    -- | An empty annotation environment.
    emptyAnnEnv :: AnnEnv
    emptyAnnEnv = MkAnnEnv emptyModuleEnv emptyNameEnv

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

    -- | Add the given annotation to the environment.
    extendAnnEnvList :: AnnEnv -> [Annotation] -> AnnEnv
    extendAnnEnvList env =
      foldl' extendAnnEnv env

    extendAnnEnv :: AnnEnv -> Annotation -> AnnEnv
    extendAnnEnv (MkAnnEnv mod_env name_env) (Annotation tgt payload) =
      case tgt of
        NamedTarget name -> MkAnnEnv mod_env (extendNameEnv_C (++) name_env name [payload])
        ModuleTarget mod -> MkAnnEnv (extendModuleEnvWith (++) mod_env mod [payload]) name_env

    -- | Union two annotation environments.
    plusAnnEnv :: AnnEnv -> AnnEnv -> AnnEnv
    plusAnnEnv a b =
      MkAnnEnv { ann_mod_env = plusModuleEnv_C (++) (ann_mod_env a) (ann_mod_env b)
               , ann_name_env = plusNameEnv_C (++) (ann_name_env a) (ann_name_env 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 deserialize env
      = mapMaybe (fromSerialized deserialize) . findAnnPayloads 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 env target tyrep
      = [ ws | Serialized tyrep' ws <- findAnnPayloads env target
        , tyrep' == tyrep ]

    -- | Find payloads for the given 'CoreAnnTarget' in an 'AnnEnv'.
    findAnnPayloads :: AnnEnv -> CoreAnnTarget -> [AnnPayload]
    findAnnPayloads env target =
      case target of
        ModuleTarget mod -> lookupWithDefaultModuleEnv (ann_mod_env env) [] mod
        NamedTarget name -> fromMaybe [] $ lookupNameEnv (ann_name_env env) 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 deserialize env
      = ( mapModuleEnv deserAnns (ann_mod_env env)
        , mapNameEnv deserAnns (ann_name_env env)
        )
      where deserAnns = mapMaybe (fromSerialized deserialize)