module GHC.Types.Annotations (
Annotation(..), AnnPayload,
AnnTarget(..), CoreAnnTarget,
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 )
data Annotation = Annotation {
ann_target :: CoreAnnTarget,
ann_value :: AnnPayload
}
type AnnPayload = Serialized
data AnnTarget name
= NamedTarget name
| ModuleTarget Module
deriving (Functor)
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)
data AnnEnv = MkAnnEnv { ann_mod_env :: !(ModuleEnv [AnnPayload])
, ann_name_env :: !(NameEnv [AnnPayload])
}
emptyAnnEnv :: AnnEnv
emptyAnnEnv = MkAnnEnv emptyModuleEnv emptyNameEnv
mkAnnEnv :: [Annotation] -> AnnEnv
mkAnnEnv = extendAnnEnvList emptyAnnEnv
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
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)
}
findAnns :: Typeable a => ([Word8] -> a) -> AnnEnv -> CoreAnnTarget -> [a]
findAnns deserialize env
= mapMaybe (fromSerialized deserialize) . findAnnPayloads env
findAnnsByTypeRep :: AnnEnv -> CoreAnnTarget -> TypeRep -> [[Word8]]
findAnnsByTypeRep env target tyrep
= [ ws | Serialized tyrep' ws <- findAnnPayloads env target
, tyrep' == tyrep ]
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
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)