-- |
-- 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 Annotations (
        -- * Main Annotation data types
        Annotation(..), AnnPayload,
        AnnTarget(..), CoreAnnTarget,
        getAnnTargetName_maybe,

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

import GhcPrelude

import Binary
import Module           ( Module )
import Name
import Outputable
import GHC.Serialized
import UniqFM
import Unique

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 'HsDecls.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 (a -> AnnTarget b -> AnnTarget a
(a -> b) -> AnnTarget a -> AnnTarget b
(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
<$ :: a -> AnnTarget b -> AnnTarget a
$c<$ :: forall a b. a -> AnnTarget b -> AnnTarget a
fmap :: (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

-- | Get the 'name' of an annotation target if it exists.
getAnnTargetName_maybe :: AnnTarget name -> Maybe name
getAnnTargetName_maybe :: AnnTarget name -> Maybe name
getAnnTargetName_maybe (NamedTarget name
nm) = name -> Maybe name
forall a. a -> Maybe a
Just name
nm
getAnnTargetName_maybe AnnTarget name
_                = Maybe name
forall a. Maybe a
Nothing

instance Uniquable name => Uniquable (AnnTarget name) where
    getUnique :: AnnTarget name -> Unique
getUnique (NamedTarget name
nm) = name -> Unique
forall a. Uniquable a => a -> Unique
getUnique name
nm
    getUnique (ModuleTarget Module
mod) = Unique -> Int -> Unique
deriveUnique (Module -> Unique
forall a. Uniquable a => a -> Unique
getUnique Module
mod) Int
0
    -- deriveUnique prevents OccName uniques clashing with NamedTarget

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
-- Can't use a type synonym or we hit bug #2412 due to source import
newtype AnnEnv = MkAnnEnv (UniqFM [AnnPayload])

-- | An empty annotation environment.
emptyAnnEnv :: AnnEnv
emptyAnnEnv :: AnnEnv
emptyAnnEnv = UniqFM [AnnPayload] -> AnnEnv
MkAnnEnv UniqFM [AnnPayload]
forall elt. UniqFM elt
emptyUFM

-- | 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 (MkAnnEnv UniqFM [AnnPayload]
env) [Annotation]
anns
  = UniqFM [AnnPayload] -> AnnEnv
MkAnnEnv (UniqFM [AnnPayload] -> AnnEnv) -> UniqFM [AnnPayload] -> AnnEnv
forall a b. (a -> b) -> a -> b
$ ([AnnPayload] -> [AnnPayload] -> [AnnPayload])
-> UniqFM [AnnPayload]
-> [(Unique, [AnnPayload])]
-> UniqFM [AnnPayload]
forall key elt.
Uniquable key =>
(elt -> elt -> elt) -> UniqFM elt -> [(key, elt)] -> UniqFM elt
addListToUFM_C [AnnPayload] -> [AnnPayload] -> [AnnPayload]
forall a. [a] -> [a] -> [a]
(++) UniqFM [AnnPayload]
env ([(Unique, [AnnPayload])] -> UniqFM [AnnPayload])
-> [(Unique, [AnnPayload])] -> UniqFM [AnnPayload]
forall a b. (a -> b) -> a -> b
$
    (Annotation -> (Unique, [AnnPayload]))
-> [Annotation] -> [(Unique, [AnnPayload])]
forall a b. (a -> b) -> [a] -> [b]
map (\Annotation
ann -> (CoreAnnTarget -> Unique
forall a. Uniquable a => a -> Unique
getUnique (Annotation -> CoreAnnTarget
ann_target Annotation
ann), [Annotation -> AnnPayload
ann_value Annotation
ann])) [Annotation]
anns

-- | Union two annotation environments.
plusAnnEnv :: AnnEnv -> AnnEnv -> AnnEnv
plusAnnEnv :: AnnEnv -> AnnEnv -> AnnEnv
plusAnnEnv (MkAnnEnv UniqFM [AnnPayload]
env1) (MkAnnEnv UniqFM [AnnPayload]
env2) = UniqFM [AnnPayload] -> AnnEnv
MkAnnEnv (UniqFM [AnnPayload] -> AnnEnv) -> UniqFM [AnnPayload] -> AnnEnv
forall a b. (a -> b) -> a -> b
$ ([AnnPayload] -> [AnnPayload] -> [AnnPayload])
-> UniqFM [AnnPayload]
-> UniqFM [AnnPayload]
-> UniqFM [AnnPayload]
forall elt.
(elt -> elt -> elt) -> UniqFM elt -> UniqFM elt -> UniqFM elt
plusUFM_C [AnnPayload] -> [AnnPayload] -> [AnnPayload]
forall a. [a] -> [a] -> [a]
(++) UniqFM [AnnPayload]
env1 UniqFM [AnnPayload]
env2

-- | 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 :: ([Word8] -> a) -> AnnEnv -> CoreAnnTarget -> [a]
findAnns [Word8] -> a
deserialize (MkAnnEnv UniqFM [AnnPayload]
ann_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
. (UniqFM [AnnPayload]
-> [AnnPayload] -> CoreAnnTarget -> [AnnPayload]
forall key elt. Uniquable key => UniqFM elt -> elt -> key -> elt
lookupWithDefaultUFM UniqFM [AnnPayload]
ann_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 (MkAnnEnv UniqFM [AnnPayload]
ann_env) CoreAnnTarget
target TypeRep
tyrep
  = [ [Word8]
ws | Serialized TypeRep
tyrep' [Word8]
ws <- UniqFM [AnnPayload]
-> [AnnPayload] -> CoreAnnTarget -> [AnnPayload]
forall key elt. Uniquable key => UniqFM elt -> elt -> key -> elt
lookupWithDefaultUFM UniqFM [AnnPayload]
ann_env [] CoreAnnTarget
target
    , TypeRep
tyrep' TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
== TypeRep
tyrep ]

-- | 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 -> UniqFM [a]
deserializeAnns :: ([Word8] -> a) -> AnnEnv -> UniqFM [a]
deserializeAnns [Word8] -> a
deserialize (MkAnnEnv UniqFM [AnnPayload]
ann_env)
  = ([AnnPayload] -> [a]) -> UniqFM [AnnPayload] -> UniqFM [a]
forall elt1 elt2. (elt1 -> elt2) -> UniqFM elt1 -> UniqFM elt2
mapUFM ((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)) UniqFM [AnnPayload]
ann_env