Give Uniq[D]FM a phantom type for its key.

This fixes #17667 and should help to avoid such issues going forward.

The changes are mostly mechanical in nature. With two notable
exceptions.

* The register allocator.

  The register allocator references registers by distinct uniques.
  However they come from the types of VirtualReg, Reg or Unique in
  various places. As a result we sometimes cast the key type of the
  map and use functions which operate on the now typed map but take
  a raw Unique as actual key. The logic itself has not changed it
  just becomes obvious where we do so now.

* <Type>Env Modules.

As an example a ClassEnv is currently queried using the types `Class`,
`Name`, and `TyCon`. This is safe since for a distinct class value all
these expressions give the same unique.

    getUnique cls
    getUnique (classTyCon cls)
    getUnique (className cls)
    getUnique (tcName $ classTyCon cls)

This is for the most part contained within the modules defining the
interface. However it requires us to play dirty when we are given a
`Name` to lookup in a `UniqFM Class a` map. But again the logic did
not change and it's for the most part hidden behind the Env Module.

Some of these cases could be avoided by refactoring but this is left
for future work.

We also bump the haddock submodule as it uses UniqFM.

16 files changed, 176 insertions, 86 deletions

diff --git a/compiler/GHC/CmmToAsm/BlockLayout.hs b/compiler/GHC/CmmToAsm/BlockLayout.hs
index 08af8e9f9f..57d265782b 100644
--- a/compiler/GHC/CmmToAsm/BlockLayout.hs
+++ b/compiler/GHC/CmmToAsm/BlockLayout.hs
@@ -29,7 +29,6 @@ import GHC.Platform
 import GHC.Driver.Session (gopt, GeneralFlag(..), DynFlags, targetPlatform)
 import GHC.Types.Unique.FM
 import GHC.Utils.Misc
-import GHC.Types.Unique
 
 import GHC.Data.Graph.Directed
 import GHC.Utils.Outputable
@@ -926,8 +925,8 @@ seqBlocks infos blocks = placeNext pullable0 todo0
         = pprPanic "seqBlocks" (ppr tooManyNextNodes)
 
 
-lookupDeleteUFM :: Uniquable key => UniqFM elt -> key
-                -> Maybe (elt, UniqFM elt)
+lookupDeleteUFM :: UniqFM BlockId elt -> BlockId
+                -> Maybe (elt, UniqFM BlockId elt)
 lookupDeleteUFM m k = do -- Maybe monad
     v <- lookupUFM m k
     return (v, delFromUFM m k)
diff --git a/compiler/GHC/CmmToAsm/Monad.hs b/compiler/GHC/CmmToAsm/Monad.hs
index 88d8f4b17c..2827e7026c 100644
--- a/compiler/GHC/CmmToAsm/Monad.hs
+++ b/compiler/GHC/CmmToAsm/Monad.hs
@@ -119,7 +119,7 @@ data NatM_State
         -- generated instructions. So instead we update the CFG as we go.
         }
 
-type DwarfFiles = UniqFM (FastString, Int)
+type DwarfFiles = UniqFM FastString (FastString, Int)
 
 newtype NatM result = NatM (NatM_State -> (result, NatM_State))
     deriving (Functor)
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph.hs b/compiler/GHC/CmmToAsm/Reg/Graph.hs
index 022c9eed4c..fad2750ef4 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph.hs
@@ -46,7 +46,7 @@ maxSpinCount    = 10
 regAlloc
         :: (Outputable statics, Outputable instr, Instruction instr)
         => NCGConfig
-        -> UniqFM (UniqSet RealReg)     -- ^ registers we can use for allocation
+        -> UniqFM RegClass (UniqSet RealReg)     -- ^ registers we can use for allocation
         -> UniqSet Int                  -- ^ set of available spill slots.
         -> Int                          -- ^ current number of spill slots
         -> [LiveCmmDecl statics instr]  -- ^ code annotated with liveness information.
@@ -96,7 +96,7 @@ regAlloc_spin
         -> Color.Triv VirtualReg RegClass RealReg
                 -- ^ Function for calculating whether a register is trivially
                 --   colourable.
-        -> UniqFM (UniqSet RealReg)      -- ^ Free registers that we can allocate.
+        -> UniqFM RegClass (UniqSet RealReg)      -- ^ Free registers that we can allocate.
         -> UniqSet Int                   -- ^ Free stack slots that we can use.
         -> Int                           -- ^ Number of spill slots in use
         -> [RegAllocStats statics instr] -- ^ Current regalloc stats to add to.
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/Coalesce.hs b/compiler/GHC/CmmToAsm/Reg/Graph/Coalesce.hs
index 0bdee541ed..ccf92baaf9 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/Coalesce.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/Coalesce.hs
@@ -44,7 +44,7 @@ regCoalesce code
 -- | Add a v1 = v2 register renaming to the map.
 --   The register with the lowest lexical name is set as the
 --   canonical version.
-buildAlloc :: UniqFM Reg -> (Reg, Reg) -> UniqFM Reg
+buildAlloc :: UniqFM Reg Reg -> (Reg, Reg) -> UniqFM Reg Reg
 buildAlloc fm (r1, r2)
  = let  rmin    = min r1 r2
         rmax    = max r1 r2
@@ -53,7 +53,7 @@ buildAlloc fm (r1, r2)
 
 -- | Determine the canonical name for a register by following
 --   v1 = v2 renamings in this map.
-sinkReg :: UniqFM Reg -> Reg -> Reg
+sinkReg :: UniqFM Reg Reg -> Reg -> Reg
 sinkReg fm r
  = case lookupUFM fm r of
         Nothing -> r
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs b/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
index 4694ba6b96..6226c1c269 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/Spill.hs
@@ -10,6 +10,7 @@ module GHC.CmmToAsm.Reg.Graph.Spill (
 import GHC.Prelude
 
 import GHC.CmmToAsm.Reg.Liveness
+import GHC.CmmToAsm.Reg.Utils
 import GHC.CmmToAsm.Instr
 import GHC.Platform.Reg
 import GHC.Cmm hiding (RegSet)
@@ -69,8 +70,11 @@ regSpill platform code slotsFree slotCount regs
         = do
                 -- Allocate a slot for each of the spilled regs.
                 let slots       = take (sizeUniqSet regs) $ nonDetEltsUniqSet slotsFree
-                let regSlotMap  = listToUFM
-                                $ zip (nonDetEltsUniqSet regs) slots
+                let
+                    regSlotMap  = toRegMap -- Cast keys from VirtualReg to Reg
+                                           -- See Note [UniqFM and the register allocator]
+                                $ listToUFM
+                                $ zip (nonDetEltsUniqSet regs) slots :: UniqFM Reg Int
                     -- This is non-deterministic but we do not
                     -- currently support deterministic code-generation.
                     -- See Note [Unique Determinism and code generation]
@@ -158,7 +162,7 @@ regSpill_top platform regSlotMap cmm
 regSpill_block
         :: Instruction instr
         => Platform
-        -> UniqFM Int   -- ^ map of vregs to slots they're being spilled to.
+        -> UniqFM Reg Int   -- ^ map of vregs to slots they're being spilled to.
         -> LiveBasicBlock instr
         -> SpillM (LiveBasicBlock instr)
 
@@ -174,7 +178,7 @@ regSpill_block platform regSlotMap (BasicBlock i instrs)
 regSpill_instr
         :: Instruction instr
         => Platform
-        -> UniqFM Int -- ^ map of vregs to slots they're being spilled to.
+        -> UniqFM Reg Int -- ^ map of vregs to slots they're being spilled to.
         -> LiveInstr instr
         -> SpillM [LiveInstr instr]
 
@@ -223,7 +227,7 @@ regSpill_instr platform regSlotMap
 --   writes to a vreg that is being spilled.
 spillRead
         :: Instruction instr
-        => UniqFM Int
+        => UniqFM Reg Int
         -> instr
         -> Reg
         -> SpillM (instr, ([LiveInstr instr'], [LiveInstr instr']))
@@ -246,7 +250,7 @@ spillRead regSlotMap instr reg
 --   writes to a vreg that is being spilled.
 spillWrite
         :: Instruction instr
-        => UniqFM Int
+        => UniqFM Reg Int
         -> instr
         -> Reg
         -> SpillM (instr, ([LiveInstr instr'], [LiveInstr instr']))
@@ -269,7 +273,7 @@ spillWrite regSlotMap instr reg
 --   both reads and writes to a vreg that is being spilled.
 spillModify
         :: Instruction instr
-        => UniqFM Int
+        => UniqFM Reg Int
         -> instr
         -> Reg
         -> SpillM (instr, ([LiveInstr instr'], [LiveInstr instr']))
@@ -334,7 +338,7 @@ data SpillS
           stateUS       :: UniqSupply
 
           -- | Spilled vreg vs the number of times it was loaded, stored.
-        , stateSpillSL  :: UniqFM (Reg, Int, Int) }
+        , stateSpillSL  :: UniqFM Reg (Reg, Int, Int) }
 
 
 -- | Create a new spiller state.
@@ -366,7 +370,7 @@ accSpillSL (r1, s1, l1) (_, s2, l2)
 --   Tells us what registers were spilled.
 data SpillStats
         = SpillStats
-        { spillStoreLoad        :: UniqFM (Reg, Int, Int) }
+        { spillStoreLoad        :: UniqFM Reg (Reg, Int, Int) }
 
 
 -- | Extract spiller statistics from the spiller state.
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs b/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
index a5016abc6f..fec35cb6bc 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/SpillClean.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 -- | Clean out unneeded spill\/reload instructions.
 --
@@ -340,7 +341,7 @@ cleanBackward liveSlotsOnEntry noReloads acc lis
 cleanBackward'
         :: Instruction instr
         => BlockMap IntSet
-        -> UniqFM [BlockId]
+        -> UniqFM Store [BlockId]
         -> UniqSet Int
         -> [LiveInstr instr]
         -> [LiveInstr instr]
@@ -438,17 +439,17 @@ type CleanM
 data CleanS
         = CleanS
         { -- | Regs which are valid at the start of each block.
-          sJumpValid            :: UniqFM (Assoc Store)
+          sJumpValid            :: UniqFM BlockId (Assoc Store)
 
           -- | Collecting up what regs were valid across each jump.
           --    in the next pass we can collate these and write the results
           --    to sJumpValid.
-        , sJumpValidAcc         :: UniqFM [Assoc Store]
+        , sJumpValidAcc         :: UniqFM BlockId [Assoc Store]
 
           -- | Map of (slot -> blocks which reload from this slot)
           --    used to decide if whether slot spilled to will ever be
           --    reloaded from on this path.
-        , sReloadedBy           :: UniqFM [BlockId]
+        , sReloadedBy           :: UniqFM Store [BlockId]
 
           -- | Spills and reloads cleaned each pass (latest at front)
         , sCleanedCount         :: [(Int, Int)]
@@ -533,7 +534,8 @@ instance Outputable Store where
 -- In the spill cleaner, two store locations are associated if they are known
 -- to hold the same value.
 --
-type Assoc a    = UniqFM (UniqSet a)
+-- TODO: Monomorphize: I think we only ever use this with a ~ Store
+type Assoc a    = UniqFM a (UniqSet a)
 
 -- | An empty association
 emptyAssoc :: Assoc a
@@ -541,8 +543,9 @@ emptyAssoc      = emptyUFM
 
 
 -- | Add an association between these two things.
-addAssoc :: Uniquable a
-         => a -> a -> Assoc a -> Assoc a
+-- addAssoc :: Uniquable a
+--          => a -> a -> Assoc a -> Assoc a
+addAssoc :: Store -> Store -> Assoc Store -> Assoc Store
 
 addAssoc a b m
  = let  m1      = addToUFM_C unionUniqSets m  a (unitUniqSet b)
@@ -551,9 +554,7 @@ addAssoc a b m
 
 
 -- | Delete all associations to a node.
-delAssoc :: (Uniquable a)
-         => a -> Assoc a -> Assoc a
-
+delAssoc :: Store -> Assoc Store -> Assoc Store
 delAssoc a m
         | Just aSet     <- lookupUFM  m a
         , m1            <- delFromUFM m a
@@ -565,9 +566,7 @@ delAssoc a m
 
 
 -- | Delete a single association edge (a -> b).
-delAssoc1 :: Uniquable a
-          => a -> a -> Assoc a -> Assoc a
-
+delAssoc1 :: Store -> Store -> Assoc Store -> Assoc Store
 delAssoc1 a b m
         | Just aSet     <- lookupUFM m a
         = addToUFM m a (delOneFromUniqSet aSet b)
@@ -576,17 +575,14 @@ delAssoc1 a b m
 
 
 -- | Check if these two things are associated.
-elemAssoc :: (Uniquable a)
-          => a -> a -> Assoc a -> Bool
+elemAssoc :: Store -> Store -> Assoc Store -> Bool
 
 elemAssoc a b m
         = elementOfUniqSet b (closeAssoc a m)
 
 
 -- | Find the refl. trans. closure of the association from this point.
-closeAssoc :: (Uniquable a)
-        => a -> Assoc a -> UniqSet a
-
+closeAssoc :: Store -> Assoc Store -> UniqSet Store
 closeAssoc a assoc
  =      closeAssoc' assoc emptyUniqSet (unitUniqSet a)
  where
@@ -615,6 +611,6 @@ closeAssoc a assoc
                         (unionUniqSets   toVisit neighbors)
 
 -- | Intersect two associations.
-intersectAssoc :: Assoc a -> Assoc a -> Assoc a
+intersectAssoc :: Assoc Store -> Assoc Store -> Assoc Store
 intersectAssoc a b
         = intersectUFM_C (intersectUniqSets) a b
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs b/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
index 995b286839..1ea380fabf 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/SpillCost.hs
@@ -48,9 +48,9 @@ type SpillCostRecord
 
 -- | Map of `SpillCostRecord`
 type SpillCostInfo
-        = UniqFM SpillCostRecord
+        = UniqFM VirtualReg SpillCostRecord
 
-type SpillCostState = State (UniqFM SpillCostRecord) ()
+type SpillCostState = State SpillCostInfo ()
 
 -- | An empty map of spill costs.
 zeroSpillCostInfo :: SpillCostInfo
@@ -264,7 +264,7 @@ spillCost_length info _ reg
 
 
 -- | Extract a map of register lifetimes from a `SpillCostInfo`.
-lifeMapFromSpillCostInfo :: SpillCostInfo -> UniqFM (VirtualReg, Int)
+lifeMapFromSpillCostInfo :: SpillCostInfo -> UniqFM VirtualReg (VirtualReg, Int)
 lifeMapFromSpillCostInfo info
         = listToUFM
         $ map (\(r, _, _, life) -> (r, (r, life)))
diff --git a/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs b/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs
index ddd353c4f2..a0b1519a93 100644
--- a/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Graph/Stats.hs
@@ -64,7 +64,7 @@ data RegAllocStats statics instr
         , raGraph       :: Color.Graph VirtualReg RegClass RealReg
 
           -- | The regs that were coalesced.
-        , raCoalesced   :: UniqFM VirtualReg
+        , raCoalesced   :: UniqFM VirtualReg VirtualReg
 
           -- | Spiller stats.
         , raSpillStats  :: SpillStats
@@ -88,7 +88,7 @@ data RegAllocStats statics instr
         , raGraphColored  :: Color.Graph VirtualReg RegClass RealReg
 
           -- | Regs that were coalesced.
-        , raCoalesced     :: UniqFM VirtualReg
+        , raCoalesced     :: UniqFM VirtualReg VirtualReg
 
           -- | Code with coalescings applied.
         , raCodeCoalesced :: [LiveCmmDecl statics instr]
@@ -242,7 +242,7 @@ pprStatsLifetimes stats
         $$ text "\n")
 
 
-binLifetimeCount :: UniqFM (VirtualReg, Int) -> UniqFM (Int, Int)
+binLifetimeCount :: UniqFM VirtualReg (VirtualReg, Int) -> UniqFM Int (Int, Int)
 binLifetimeCount fm
  = let  lifes   = map (\l -> (l, (l, 1)))
                 $ map snd
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear.hs b/compiler/GHC/CmmToAsm/Reg/Linear.hs
index 55cb73af1a..f777a21ca6 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear.hs
@@ -119,6 +119,7 @@ import qualified GHC.CmmToAsm.Reg.Linear.X86    as X86
 import qualified GHC.CmmToAsm.Reg.Linear.X86_64 as X86_64
 import GHC.CmmToAsm.Reg.Target
 import GHC.CmmToAsm.Reg.Liveness
+import GHC.CmmToAsm.Reg.Utils
 import GHC.CmmToAsm.Instr
 import GHC.CmmToAsm.Config
 import GHC.Platform.Reg
@@ -427,7 +428,7 @@ raInsn _     new_instrs _ (LiveInstr ii@(Instr i) Nothing)
 
 raInsn block_live new_instrs id (LiveInstr (Instr instr) (Just live))
  = do
-    assig    <- getAssigR
+    assig    <- getAssigR :: RegM freeRegs (UniqFM Reg Loc)
 
     -- If we have a reg->reg move between virtual registers, where the
     -- src register is not live after this instruction, and the dst
@@ -486,7 +487,8 @@ isInReg src assig | Just (InReg _) <- lookupUFM assig src = True
                   | otherwise = False
 
 
-genRaInsn :: OutputableRegConstraint freeRegs instr
+genRaInsn :: forall freeRegs instr.
+             OutputableRegConstraint freeRegs instr
           => BlockMap RegSet
           -> [instr]
           -> BlockId
@@ -500,13 +502,13 @@ genRaInsn block_live new_instrs block_id instr r_dying w_dying = do
   platform <- getPlatform
   case regUsageOfInstr platform instr of { RU read written ->
     do
-    let real_written    = [ rr  | (RegReal     rr) <- written ]
+    let real_written    = [ rr  | (RegReal     rr) <- written ] :: [RealReg]
     let virt_written    = [ vr  | (RegVirtual  vr) <- written ]
 
     -- we don't need to do anything with real registers that are
     -- only read by this instr.  (the list is typically ~2 elements,
     -- so using nub isn't a problem).
-    let virt_read       = nub [ vr      | (RegVirtual vr) <- read ]
+    let virt_read       = nub [ vr      | (RegVirtual vr) <- read ] :: [VirtualReg]
 
     -- debugging
 {-    freeregs <- getFreeRegsR
@@ -560,15 +562,20 @@ genRaInsn block_live new_instrs block_id instr r_dying w_dying = do
 
     let
         -- (i) Patch the instruction
+        patch_map :: UniqFM Reg Reg
         patch_map
-                = listToUFM
+                = toRegMap $ -- Cast key from VirtualReg to Reg
+                             -- See Note [UniqFM and the register allocator]
+                  listToUFM
                         [ (t, RegReal r)
                                 | (t, r) <- zip virt_read    r_allocd
                                          ++ zip virt_written w_allocd ]
 
+        patched_instr :: instr
         patched_instr
                 = patchRegsOfInstr adjusted_instr patchLookup
 
+        patchLookup :: Reg -> Reg
         patchLookup x
                 = case lookupUFM patch_map x of
                         Nothing -> x
@@ -631,7 +638,8 @@ releaseRegs regs = do
 --
 
 saveClobberedTemps
-        :: (Instruction instr, FR freeRegs)
+        :: forall instr freeRegs.
+           (Instruction instr, FR freeRegs)
         => [RealReg]            -- real registers clobbered by this instruction
         -> [Reg]                -- registers which are no longer live after this insn
         -> RegM freeRegs [instr]         -- return: instructions to spill any temps that will
@@ -642,8 +650,10 @@ saveClobberedTemps [] _
 
 saveClobberedTemps clobbered dying
  = do
-        assig   <- getAssigR
-        let to_spill
+        assig   <- getAssigR :: RegM freeRegs (UniqFM Reg Loc)
+        -- Unique represents the VirtualReg
+        let to_spill :: [(Unique, RealReg)]
+            to_spill
                 = [ (temp,reg)
                         | (temp, InReg reg) <- nonDetUFMToList assig
                         -- This is non-deterministic but we do not
@@ -657,6 +667,8 @@ saveClobberedTemps clobbered dying
         return instrs
 
    where
+     -- See Note [UniqFM and the register allocator]
+     clobber :: RegMap Loc -> [instr] -> [(Unique,RealReg)] -> RegM freeRegs ([instr], RegMap Loc)
      clobber assig instrs []
             = return (instrs, assig)
 
@@ -675,7 +687,7 @@ saveClobberedTemps clobbered dying
               (my_reg : _) -> do
                   setFreeRegsR (frAllocateReg platform my_reg freeRegs)
 
-                  let new_assign = addToUFM assig temp (InReg my_reg)
+                  let new_assign = addToUFM_Directly assig temp (InReg my_reg)
                   let instr = mkRegRegMoveInstr platform
                                   (RegReal reg) (RegReal my_reg)
 
@@ -688,7 +700,7 @@ saveClobberedTemps clobbered dying
                   -- record why this reg was spilled for profiling
                   recordSpill (SpillClobber temp)
 
-                  let new_assign  = addToUFM assig temp (InBoth reg slot)
+                  let new_assign  = addToUFM_Directly assig temp (InBoth reg slot)
 
                   clobber new_assign (spill : instrs) rest
 
@@ -719,12 +731,13 @@ clobberRegs clobbered
         -- also catches temps which were loaded up during allocation
         -- of read registers, not just those saved in saveClobberedTemps.
 
+        clobber :: RegMap Loc -> [(Unique,Loc)] -> RegMap Loc
         clobber assig []
                 = assig
 
         clobber assig ((temp, InBoth reg slot) : rest)
                 | any (realRegsAlias reg) clobbered
-                = clobber (addToUFM assig temp (InMem slot)) rest
+                = clobber (addToUFM_Directly assig temp (InMem slot)) rest
 
         clobber assig (_:rest)
                 = clobber assig rest
@@ -762,8 +775,9 @@ allocateRegsAndSpill _       _    spills alloc []
         = return (spills, reverse alloc)
 
 allocateRegsAndSpill reading keep spills alloc (r:rs)
- = do   assig <- getAssigR :: RegM freeRegs (RegMap Loc)
+ = do   assig <- toVRegMap <$> getAssigR
         -- pprTraceM "allocateRegsAndSpill:assig" (ppr (r:rs) $$ ppr assig)
+        -- See Note [UniqFM and the register allocator]
         let doSpill = allocRegsAndSpill_spill reading keep spills alloc r rs assig
         case lookupUFM assig r of
                 -- case (1a): already in a register
@@ -776,7 +790,7 @@ allocateRegsAndSpill reading keep spills alloc (r:rs)
                 -- NB2. This is why we must process written registers here, even if they
                 -- are also read by the same instruction.
                 Just (InBoth my_reg _)
-                 -> do  when (not reading) (setAssigR (addToUFM assig r (InReg my_reg)))
+                 -> do  when (not reading) (setAssigR $ toRegMap (addToUFM assig r (InReg my_reg)))
                         allocateRegsAndSpill reading keep spills (my_reg:alloc) rs
 
                 -- Not already in a register, so we need to find a free one...
@@ -801,15 +815,14 @@ allocateRegsAndSpill reading keep spills alloc (r:rs)
 -- Note: I tried returning a list of past assignments, but that
 -- turned out to barely matter but added a few tenths of
 -- a percent to compile time.
-findPrefRealReg :: forall freeRegs u. Uniquable u
-               => u -> RegM freeRegs (Maybe RealReg)
+findPrefRealReg :: VirtualReg -> RegM freeRegs (Maybe RealReg)
 findPrefRealReg vreg = do
   bassig <- getBlockAssigR :: RegM freeRegs (BlockMap (freeRegs,RegMap Loc))
   return $ foldr (findVirtRegAssig) Nothing bassig
   where
     findVirtRegAssig :: (freeRegs,RegMap Loc) -> Maybe RealReg -> Maybe RealReg
     findVirtRegAssig assig z =
-        z <|>   case lookupUFM (snd assig) vreg of
+        z <|>   case lookupUFM (toVRegMap $ snd assig) vreg of
                         Just (InReg real_reg) -> Just real_reg
                         Just (InBoth real_reg _) -> Just real_reg
                         _ -> z
@@ -823,7 +836,7 @@ allocRegsAndSpill_spill :: (FR freeRegs, Instruction instr, Outputable instr)
                         -> [RealReg]
                         -> VirtualReg
                         -> [VirtualReg]
-                        -> UniqFM Loc
+                        -> UniqFM VirtualReg Loc
                         -> SpillLoc
                         -> RegM freeRegs ([instr], [RealReg])
 allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
@@ -845,7 +858,8 @@ allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
                         = first_free
                 spills'   <- loadTemp r spill_loc final_reg spills
 
-                setAssigR       (addToUFM assig r $! newLocation spill_loc final_reg)
+                setAssigR $ toRegMap
+                          $ (addToUFM assig r $! newLocation spill_loc final_reg)
                 setFreeRegsR $  frAllocateReg platform final_reg freeRegs
 
                 allocateRegsAndSpill reading keep spills' (final_reg : alloc) rs
@@ -856,7 +870,7 @@ allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
            do   let inRegOrBoth (InReg _) = True
                     inRegOrBoth (InBoth _ _) = True
                     inRegOrBoth _ = False
-                let candidates' :: UniqFM Loc
+                let candidates' :: UniqFM VirtualReg Loc
                     candidates' =
                       flip delListFromUFM keep $
                       filterUFM inRegOrBoth $
@@ -867,7 +881,8 @@ allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
                 let candidates = nonDetUFMToList candidates'
 
                 -- the vregs we could kick out that are already in a slot
-                let candidates_inBoth
+                let candidates_inBoth :: [(Unique, RealReg, StackSlot)]
+                    candidates_inBoth
                         = [ (temp, reg, mem)
                           | (temp, InBoth reg mem) <- candidates
                           , targetClassOfRealReg platform reg == classOfVirtualReg r ]
@@ -885,10 +900,10 @@ allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
                         -- just free up its register for use.
                         | (temp, my_reg, slot) : _      <- candidates_inBoth
                         = do    spills' <- loadTemp r spill_loc my_reg spills
-                                let assig1  = addToUFM assig temp (InMem slot)
+                                let assig1  = addToUFM_Directly assig temp (InMem slot)
                                 let assig2  = addToUFM assig1 r $! newLocation spill_loc my_reg
 
-                                setAssigR assig2
+                                setAssigR $ toRegMap assig2
                                 allocateRegsAndSpill reading keep spills' (my_reg:alloc) rs
 
                         -- otherwise, we need to spill a temporary that currently
@@ -905,9 +920,9 @@ allocRegsAndSpill_spill reading keep spills alloc r rs assig spill_loc
                                 recordSpill (SpillAlloc temp_to_push_out)
 
                                 -- update the register assignment
-                                let assig1  = addToUFM assig temp_to_push_out   (InMem slot)
+                                let assig1  = addToUFM_Directly assig temp_to_push_out   (InMem slot)
                                 let assig2  = addToUFM assig1 r                 $! newLocation spill_loc my_reg
-                                setAssigR assig2
+                                setAssigR $ toRegMap assig2
 
                                 -- if need be, load up a spilled temp into the reg we've just freed up.
                                 spills' <- loadTemp r spill_loc my_reg spills
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs b/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
index 6a110f0a48..33a15fd7b8 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/Base.hs
@@ -100,7 +100,9 @@ data SpillReason
 -- | Used to carry interesting stats out of the register allocator.
 data RegAllocStats
         = RegAllocStats
-        { ra_spillInstrs        :: UniqFM [Int]
+        { ra_spillInstrs        :: UniqFM Unique [Int] -- Keys are the uniques of regs
+                                                       -- and taken from SpillReason
+                                                       -- See Note [UniqFM and the register allocator]
         , ra_fixupList     :: [(BlockId,BlockId,BlockId)]
         -- ^ (from,fixup,to) : We inserted fixup code between from and to
         }
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs b/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
index 4ceaf4573b..8d3a46f490 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/JoinToTargets.hs
@@ -146,8 +146,8 @@ joinToTargets_again :: (Instruction instr, FR freeRegs, Outputable instr)
                     -> instr
                     -> BlockId
                     -> [BlockId]
-                    -> UniqFM Loc
-                    -> UniqFM Loc
+                    -> UniqFM Reg Loc
+                    -> UniqFM Reg Loc
                     -> RegM freeRegs ([NatBasicBlock instr], instr)
 joinToTargets_again
     block_live new_blocks block_id instr dest dests
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/StackMap.hs b/compiler/GHC/CmmToAsm/Reg/Linear/StackMap.hs
index 29864f9752..97e04936b0 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/StackMap.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/StackMap.hs
@@ -34,8 +34,9 @@ data StackMap
         { -- | The slots that are still available to be allocated.
           stackMapNextFreeSlot  :: !Int
 
+          -- See Note [UniqFM and the register allocator]
           -- | Assignment of vregs to stack slots.
-        , stackMapAssignment    :: UniqFM StackSlot }
+        , stackMapAssignment    :: UniqFM Unique StackSlot }
 
 
 -- | An empty stack map, with all slots available.
diff --git a/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs b/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs
index 414128b32c..6411e5285d 100644
--- a/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Linear/Stats.hs
@@ -11,6 +11,7 @@ import GHC.Prelude
 import GHC.CmmToAsm.Reg.Linear.Base
 import GHC.CmmToAsm.Reg.Liveness
 import GHC.CmmToAsm.Instr
+import GHC.Types.Unique (Unique)
 
 import GHC.Types.Unique.FM
 import GHC.Utils.Outputable
@@ -19,8 +20,8 @@ import GHC.Utils.Monad.State
 
 -- | Build a map of how many times each reg was alloced, clobbered, loaded etc.
 binSpillReasons
-        :: [SpillReason] -> UniqFM [Int]
-
+        :: [SpillReason] -> UniqFM Unique [Int]
+        -- See Note [UniqFM and the register allocator]
 binSpillReasons reasons
         = addListToUFM_C
                 (zipWith (+))
@@ -61,6 +62,8 @@ pprStats
 
 pprStats code statss
  = let  -- sum up all the instrs inserted by the spiller
+        -- See Note [UniqFM and the register allocator]
+        spills :: UniqFM Unique [Int]
         spills          = foldl' (plusUFM_C (zipWith (+)))
                                 emptyUFM
                         $ map ra_spillInstrs statss
diff --git a/compiler/GHC/CmmToAsm/Reg/Liveness.hs b/compiler/GHC/CmmToAsm/Reg/Liveness.hs
index f650ad6186..13dbcc5f70 100644
--- a/compiler/GHC/CmmToAsm/Reg/Liveness.hs
+++ b/compiler/GHC/CmmToAsm/Reg/Liveness.hs
@@ -66,9 +66,14 @@ import Data.IntSet              (IntSet)
 -----------------------------------------------------------------------------
 type RegSet = UniqSet Reg
 
-type RegMap a = UniqFM a
+-- | Map from some kind of register to a.
+--
+-- While we give the type for keys as Reg which is the common case
+-- sometimes we end up using VirtualReq or naked Uniques.
+-- See Note [UniqFM and the register allocator]
+type RegMap a = UniqFM Reg a
 
-emptyRegMap :: UniqFM a
+emptyRegMap :: RegMap a
 emptyRegMap = emptyUFM
 
 emptyRegSet :: RegSet
@@ -76,6 +81,9 @@ emptyRegSet = emptyUniqSet
 
 type BlockMap a = LabelMap a
 
+type SlotMap a = UniqFM Slot a
+
+type Slot = Int
 
 -- | A top level thing which carries liveness information.
 type LiveCmmDecl statics instr
@@ -400,7 +408,7 @@ slurpReloadCoalesce live
            in   unionManyBags (cs : moveBags)
 
         slurpCompM :: [LiveBasicBlock instr]
-                   -> State (UniqFM [UniqFM Reg]) [Bag (Reg, Reg)]
+                   -> State (UniqFM BlockId [UniqFM Slot Reg]) [Bag (Reg, Reg)]
         slurpCompM blocks
          = do   -- run the analysis once to record the mapping across jumps.
                 mapM_   (slurpBlock False) blocks
@@ -412,7 +420,7 @@ slurpReloadCoalesce live
                 mapM    (slurpBlock True) blocks
 
         slurpBlock :: Bool -> LiveBasicBlock instr
-                   -> State (UniqFM [UniqFM Reg]) (Bag (Reg, Reg))
+                   -> State (UniqFM BlockId [UniqFM Slot Reg]) (Bag (Reg, Reg))
         slurpBlock propagate (BasicBlock blockId instrs)
          = do   -- grab the slot map for entry to this block
                 slotMap         <- if propagate
@@ -422,12 +430,12 @@ slurpReloadCoalesce live
                 (_, mMoves)     <- mapAccumLM slurpLI slotMap instrs
                 return $ listToBag $ catMaybes mMoves
 
-        slurpLI :: UniqFM Reg                           -- current slotMap
+        slurpLI :: SlotMap Reg                           -- current slotMap
                 -> LiveInstr instr
-                -> State (UniqFM [UniqFM Reg])          -- blockId -> [slot -> reg]
+                -> State (UniqFM BlockId [SlotMap Reg])  -- blockId -> [slot -> reg]
                                                         --      for tracking slotMaps across jumps
 
-                         ( UniqFM Reg                   -- new slotMap
+                         ( SlotMap Reg           -- new slotMap
                          , Maybe (Reg, Reg))            -- maybe a new coalesce edge
 
         slurpLI slotMap li
@@ -467,15 +475,18 @@ slurpReloadCoalesce live
                 let slotMaps    = fromMaybe [] (lookupUFM map blockId)
                 return          $ foldr mergeSlotMaps emptyUFM slotMaps
 
-        mergeSlotMaps :: UniqFM Reg -> UniqFM Reg -> UniqFM Reg
+        mergeSlotMaps :: SlotMap Reg -> SlotMap Reg -> SlotMap Reg
         mergeSlotMaps map1 map2
-                = listToUFM
+                -- toList sadly means we have to use the _Directly style
+                -- functions.
+                -- TODO: We shouldn't need to go through a list here.
+                = listToUFM_Directly
                 $ [ (k, r1)
                   | (k, r1) <- nonDetUFMToList map1
                   -- This is non-deterministic but we do not
                   -- currently support deterministic code-generation.
                   -- See Note [Unique Determinism and code generation]
-                  , case lookupUFM map2 k of
+                  , case lookupUFM_Directly map2 k of
                           Nothing -> False
                           Just r2 -> r1 == r2 ]
 
diff --git a/compiler/GHC/CmmToAsm/Reg/Utils.hs b/compiler/GHC/CmmToAsm/Reg/Utils.hs
new file mode 100644
index 0000000000..3a832963fe
--- /dev/null
+++ b/compiler/GHC/CmmToAsm/Reg/Utils.hs
@@ -0,0 +1,59 @@
+module GHC.CmmToAsm.Reg.Utils
+    ( toRegMap, toVRegMap )
+where
+
+{- Note [UniqFM and the register allocator]
+   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+   Before UniqFM had a key type the register allocator
+   wasn't picky about key types, using VirtualReg, Reg
+   and Unique at various use sites for the same map.
+
+   This is safe.
+   * The Unique values come from registers at various
+     points where we lose a reference to the original
+     register value, but the unique is still valid.
+
+   * VirtualReg is a subset of the registers in Reg's type.
+     Making a value of VirtualReg into a Reg in fact doesn't
+     change its unique. This is because Reg consists of virtual
+     regs and real regs, whose unique values do not overlap.
+
+   * Since the code was written in the assumption that keys are
+     not typed it's hard to reverse this assumption now. So we get
+     some gnarly but correct code where we often pass around Uniques
+     and switch between using Uniques, VirtualReg and RealReg as keys
+     of the same map. These issues were always there. But with the
+     now-typed keys they become visible. It's a classic case of not all
+     correct programs type checking.
+
+   We reduce some of the burden by providing a way to cast
+
+        UniqFM VirtualReg a
+
+   to
+
+        UniqFM Reg a
+
+    in this module. This is safe as Reg is the sum of VirtualReg and
+    RealReg. With each kind of register keeping the same unique when
+    treated as Reg.
+
+   TODO: If you take offense to this I encourage you to refactor this
+   code. I'm sure we can do with less casting of keys and direct use
+   of uniques. It might also be reasonable to just use a IntMap directly
+   instead of dealing with UniqFM at all.
+
+
+-}
+import GHC.Types.Unique.FM
+import GHC.Platform.Reg
+
+-- These should hopefully be zero cost.
+
+toRegMap :: UniqFM VirtualReg elt -> UniqFM Reg elt
+toRegMap = unsafeCastUFMKey
+
+toVRegMap :: UniqFM Reg elt -> UniqFM VirtualReg elt
+toVRegMap = unsafeCastUFMKey
+
diff --git a/compiler/GHC/CmmToAsm/X86/RegInfo.hs b/compiler/GHC/CmmToAsm/X86/RegInfo.hs
index de11279d54..3f7d50d319 100644
--- a/compiler/GHC/CmmToAsm/X86/RegInfo.hs
+++ b/compiler/GHC/CmmToAsm/X86/RegInfo.hs
@@ -38,13 +38,13 @@ regDotColor platform reg
         Just str -> text str
         _        -> panic "Register not assigned a color"
 
-regColors :: Platform -> UniqFM [Char]
+regColors :: Platform -> UniqFM RealReg [Char]
 regColors platform = listToUFM (normalRegColors platform)
 
-normalRegColors :: Platform -> [(Reg,String)]
+normalRegColors :: Platform -> [(RealReg,String)]
 normalRegColors platform =
-    zip (map regSingle [0..lastint platform]) colors
-        ++ zip (map regSingle [firstxmm..lastxmm platform]) greys
+    zip (map realRegSingle [0..lastint platform]) colors
+        ++ zip (map realRegSingle [firstxmm..lastxmm platform]) greys
   where
     -- 16 colors - enough for amd64 gp regs
     colors = ["#800000","#ff0000","#808000","#ffff00","#008000"