diff options
-rw-r--r-- | compiler/main/Annotations.hs | 62 | ||||
-rw-r--r-- | compiler/simplCore/CoreMonad.hs | 16 | ||||
-rw-r--r-- | compiler/specialise/SpecConstr.hs | 2 | ||||
-rw-r--r-- | docs/users_guide/8.12.1-notes.rst | 8 | ||||
-rw-r--r-- | testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs | 2 | ||||
-rw-r--r-- | testsuite/tests/plugins/simple-plugin/Simple/Plugin.hs | 2 |
6 files changed, 60 insertions, 32 deletions
diff --git a/compiler/main/Annotations.hs b/compiler/main/Annotations.hs index 82d80aae43..c282217d33 100644 --- a/compiler/main/Annotations.hs +++ b/compiler/main/Annotations.hs @@ -21,12 +21,14 @@ module Annotations ( import GhcPrelude import Binary -import Module ( Module ) +import Module ( Module + , ModuleEnv, emptyModuleEnv, extendModuleEnvWith + , plusModuleEnv_C, lookupWithDefaultModuleEnv + , mapModuleEnv ) +import NameEnv import Name import Outputable import GHC.Serialized -import UniqFM -import Unique import Control.Monad import Data.Maybe @@ -60,11 +62,6 @@ getAnnTargetName_maybe :: AnnTarget name -> Maybe name getAnnTargetName_maybe (NamedTarget nm) = Just nm getAnnTargetName_maybe _ = Nothing -instance Uniquable name => Uniquable (AnnTarget name) where - getUnique (NamedTarget nm) = getUnique nm - getUnique (ModuleTarget mod) = deriveUnique (getUnique mod) 0 - -- deriveUnique prevents OccName uniques clashing with NamedTarget - instance Outputable name => Outputable (AnnTarget name) where ppr (NamedTarget nm) = text "Named target" <+> ppr nm ppr (ModuleTarget mod) = text "Module target" <+> ppr mod @@ -86,12 +83,13 @@ instance Outputable Annotation where ppr ann = ppr (ann_target 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]) +data AnnEnv = MkAnnEnv { ann_mod_env :: !(ModuleEnv [AnnPayload]) + , ann_name_env :: !(NameEnv [AnnPayload]) + } -- | An empty annotation environment. emptyAnnEnv :: AnnEnv -emptyAnnEnv = MkAnnEnv emptyUFM +emptyAnnEnv = MkAnnEnv emptyModuleEnv emptyNameEnv -- | Construct a new annotation environment that contains the list of -- annotations provided. @@ -100,33 +98,51 @@ mkAnnEnv = extendAnnEnvList emptyAnnEnv -- | Add the given annotation to the environment. extendAnnEnvList :: AnnEnv -> [Annotation] -> AnnEnv -extendAnnEnvList (MkAnnEnv env) anns - = MkAnnEnv $ addListToUFM_C (++) env $ - map (\ann -> (getUnique (ann_target ann), [ann_value ann])) anns +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 (MkAnnEnv env1) (MkAnnEnv env2) = MkAnnEnv $ plusUFM_C (++) env1 env2 +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 (MkAnnEnv ann_env) - = (mapMaybe (fromSerialized deserialize)) - . (lookupWithDefaultUFM ann_env []) +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 (MkAnnEnv ann_env) target tyrep - = [ ws | Serialized tyrep' ws <- lookupWithDefaultUFM ann_env [] target +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 -> UniqFM [a] -deserializeAnns deserialize (MkAnnEnv ann_env) - = mapUFM (mapMaybe (fromSerialized deserialize)) ann_env +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) + diff --git a/compiler/simplCore/CoreMonad.hs b/compiler/simplCore/CoreMonad.hs index fde925063b..c87bd353c0 100644 --- a/compiler/simplCore/CoreMonad.hs +++ b/compiler/simplCore/CoreMonad.hs @@ -64,10 +64,12 @@ import FastString import qualified ErrUtils as Err import ErrUtils( Severity(..) ) import UniqSupply -import UniqFM ( UniqFM, mapUFM, filterUFM ) +import NameEnv ( mapNameEnv, filterNameEnv ) import MonadUtils import NameCache +import NameEnv import SrcLoc +import Data.Bifunctor ( bimap ) import Data.List import Data.Ord import Data.Dynamic @@ -733,17 +735,19 @@ getPackageFamInstEnv = do -- annotations. -- -- See Note [Annotations] -getAnnotations :: Typeable a => ([Word8] -> a) -> ModGuts -> CoreM (UniqFM [a]) +getAnnotations :: Typeable a => ([Word8] -> a) -> ModGuts -> CoreM (ModuleEnv [a], NameEnv [a]) getAnnotations deserialize guts = do hsc_env <- getHscEnv ann_env <- liftIO $ prepareAnnotations hsc_env (Just guts) return (deserializeAnns deserialize ann_env) --- | Get at most one annotation of a given type per Unique. -getFirstAnnotations :: Typeable a => ([Word8] -> a) -> ModGuts -> CoreM (UniqFM a) +-- | Get at most one annotation of a given type per annotatable item. +getFirstAnnotations :: Typeable a => ([Word8] -> a) -> ModGuts -> CoreM (ModuleEnv a, NameEnv a) getFirstAnnotations deserialize guts - = liftM (mapUFM head . filterUFM (not . null)) - $ getAnnotations deserialize guts + = bimap mod name <$> getAnnotations deserialize guts + where + mod = mapModuleEnv head . filterModuleEnv (const $ not . null) + name = mapNameEnv head . filterNameEnv (not . null) {- Note [Annotations] diff --git a/compiler/specialise/SpecConstr.hs b/compiler/specialise/SpecConstr.hs index 8ced5a87c0..56c81ea101 100644 --- a/compiler/specialise/SpecConstr.hs +++ b/compiler/specialise/SpecConstr.hs @@ -699,7 +699,7 @@ specConstrProgram guts = do dflags <- getDynFlags us <- getUniqueSupplyM - annos <- getFirstAnnotations deserializeWithData guts + (_, annos) <- getFirstAnnotations deserializeWithData guts this_mod <- getModule let binds' = reverse $ fst $ initUs us $ do -- Note [Top-level recursive groups] diff --git a/docs/users_guide/8.12.1-notes.rst b/docs/users_guide/8.12.1-notes.rst index 49c1d623ed..94979e80c4 100644 --- a/docs/users_guide/8.12.1-notes.rst +++ b/docs/users_guide/8.12.1-notes.rst @@ -38,6 +38,14 @@ Template Haskell ``ghc`` library ~~~~~~~~~~~~~~~ + - The type of the ``getAnnotations`` function has changed to better reflect + the fact that it returns two different kinds of annotations, those on + names and those on modules: :: + + getAnnotations :: Typeable a + => ([Word8] -> a) -> ModGuts + -> CoreM (ModuleEnv [a], NameEnv [a]) + ``base`` library ~~~~~~~~~~~~~~~~ diff --git a/testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs b/testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs index 55e32e5b69..ae4135d203 100644 --- a/testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs +++ b/testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs @@ -29,5 +29,5 @@ pass g = do annotationsOn :: Data a => ModGuts -> CoreBndr -> CoreM [a] annotationsOn guts bndr = do - anns <- getAnnotations deserializeWithData guts + (_, anns) <- getAnnotations deserializeWithData guts return $ lookupWithDefaultUFM anns [] (varUnique bndr) diff --git a/testsuite/tests/plugins/simple-plugin/Simple/Plugin.hs b/testsuite/tests/plugins/simple-plugin/Simple/Plugin.hs index 938d23586c..aabc1e5b6c 100644 --- a/testsuite/tests/plugins/simple-plugin/Simple/Plugin.hs +++ b/testsuite/tests/plugins/simple-plugin/Simple/Plugin.hs @@ -46,7 +46,7 @@ findNameBndr target b mainPass :: ModGuts -> CoreM ModGuts mainPass guts = do putMsgS "Simple Plugin Pass Run" - anns <- getAnnotations deserializeWithData guts + (_, anns) <- getAnnotations deserializeWithData guts bindsOnlyPass (mapM (changeBind anns Nothing)) guts changeBind :: UniqFM [ReplaceWith] -> Maybe String -> CoreBind -> CoreM CoreBind |