content: -- |
  -- 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 )
  import Control.DeepSeq
  
  -- | 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 NFData name => NFData (AnnTarget name) where
    rnf (NamedTarget n) = rnf n
    rnf (ModuleTarget m) = rnf m
  
  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)
  
  