summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--compiler/main/Annotations.hs62
-rw-r--r--compiler/simplCore/CoreMonad.hs16
-rw-r--r--compiler/specialise/SpecConstr.hs2
-rw-r--r--docs/users_guide/8.12.1-notes.rst8
-rw-r--r--testsuite/tests/plugins/annotation-plugin/SayAnnNames.hs2
-rw-r--r--testsuite/tests/plugins/simple-plugin/Simple/Plugin.hs2
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