Drop Uniquable constraint for AnnTarget

This relied on deriveUnique, which was far too subtle to be safely
applied. Thankfully the instance doesn't appear to be used so let's just
drop it.

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