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.

4 files changed, 26 insertions, 17 deletions

diff --git a/compiler/GHC/Data/Graph/Base.hs b/compiler/GHC/Data/Graph/Base.hs
index 3c40645660..9e8cc383a4 100644
--- a/compiler/GHC/Data/Graph/Base.hs
+++ b/compiler/GHC/Data/Graph/Base.hs
@@ -45,7 +45,7 @@ type Triv k cls color
 data Graph k cls color
         = Graph {
         -- | All active nodes in the graph.
-          graphMap              :: UniqFM (Node k cls color)  }
+          graphMap              :: UniqFM k (Node k cls color)  }
 
 
 -- | An empty graph.
@@ -57,7 +57,7 @@ initGraph
 
 -- | Modify the finite map holding the nodes in the graph.
 graphMapModify
-        :: (UniqFM (Node k cls color) -> UniqFM (Node k cls color))
+        :: (UniqFM k (Node k cls color) -> UniqFM k (Node k cls color))
         -> Graph k cls color -> Graph k cls color
 
 graphMapModify f graph
diff --git a/compiler/GHC/Data/Graph/Color.hs b/compiler/GHC/Data/Graph/Color.hs
index 948447da58..e4189acb61 100644
--- a/compiler/GHC/Data/Graph/Color.hs
+++ b/compiler/GHC/Data/Graph/Color.hs
@@ -4,6 +4,7 @@
 --
 
 {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 module GHC.Data.Graph.Color (
         module GHC.Data.Graph.Base,
@@ -37,19 +38,20 @@ import Data.List
 --      the stack (ie in reverse order) and assigning them colors different to their neighbors.
 --
 colorGraph
-        :: ( Uniquable  k, Uniquable cls,  Uniquable  color
+        :: forall k cls color.
+           ( Uniquable  k, Uniquable cls,  Uniquable  color
            , Eq cls, Ord k
            , Outputable k, Outputable cls, Outputable color)
         => Bool                         -- ^ whether to do iterative coalescing
         -> Int                          -- ^ how many times we've tried to color this graph so far.
-        -> UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
+        -> UniqFM cls (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
         -> Triv   k cls color           -- ^ fn to decide whether a node is trivially colorable.
         -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
         -> Graph  k cls color           -- ^ the graph to color.
 
         -> ( Graph k cls color          -- the colored graph.
            , UniqSet k                  -- the set of nodes that we couldn't find a color for.
-           , UniqFM  k )                -- map of regs (r1 -> r2) that were coalesced
+           , UniqFM k k )                -- map of regs (r1 -> r2) that were coalesced
                                         --       r1 should be replaced by r2 in the source
 
 colorGraph iterative spinCount colors triv spill graph0
@@ -71,7 +73,7 @@ colorGraph iterative spinCount colors triv spill graph0
 
         -- run the scanner to slurp out all the trivially colorable nodes
         --      (and do coalescing if iterative coalescing is enabled)
-        (ksTriv, ksProblems, kksCoalesce2)
+        (ksTriv, ksProblems, kksCoalesce2 :: [(k,k)])
                 = colorScan iterative triv spill graph_coalesced
 
         -- If iterative coalescing is enabled, the scanner will coalesce the graph as does its business.
@@ -253,9 +255,10 @@ colorScan_spill iterative triv spill graph
 -- | Try to assign a color to all these nodes.
 
 assignColors
-        :: ( Uniquable k, Uniquable cls, Uniquable color
+        :: forall k cls color.
+           ( Uniquable k, Uniquable cls, Uniquable color
            , Outputable cls)
-        => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
+        => UniqFM cls (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
         -> Graph k cls color            -- ^ the graph
         -> [k]                          -- ^ nodes to assign a color to.
         -> ( Graph k cls color          -- the colored graph
@@ -264,7 +267,13 @@ assignColors
 assignColors colors graph ks
         = assignColors' colors graph [] ks
 
- where  assignColors' _ graph prob []
+ where  assignColors' :: UniqFM cls (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
+                        -> Graph k cls color            -- ^ the graph
+                        -> [k]                          -- ^ nodes to assign a color to.
+                        -> [k]
+                        -> ( Graph k cls color          -- the colored graph
+                        , [k])
+        assignColors' _ graph prob []
                 = (graph, prob)
 
         assignColors' colors graph prob (k:ks)
@@ -293,7 +302,7 @@ assignColors colors graph ks
 selectColor
         :: ( Uniquable k, Uniquable cls, Uniquable color
            , Outputable cls)
-        => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
+        => UniqFM cls (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
         -> Graph k cls color            -- ^ the graph
         -> k                            -- ^ key of the node to select a color for.
         -> Maybe color
diff --git a/compiler/GHC/Data/Graph/Directed.hs b/compiler/GHC/Data/Graph/Directed.hs
index c3f397051a..5bd08b9641 100644
--- a/compiler/GHC/Data/Graph/Directed.hs
+++ b/compiler/GHC/Data/Graph/Directed.hs
@@ -507,8 +507,8 @@ classifyEdges root getSucc edges =
         endFrom   = getTime ends   from
         endTo     = getTime ends   to
 
-    addTimes :: (Time, UniqFM Time, UniqFM Time) -> key
-             -> (Time, UniqFM Time, UniqFM Time)
+    addTimes :: (Time, UniqFM key Time, UniqFM key Time) -> key
+             -> (Time, UniqFM key Time, UniqFM key Time)
     addTimes (time,starts,ends) n
       --Dont reenter nodes
       | elemUFM n starts
diff --git a/compiler/GHC/Data/Graph/Ops.hs b/compiler/GHC/Data/Graph/Ops.hs
index 61f8bfe431..99e4a7eea0 100644
--- a/compiler/GHC/Data/Graph/Ops.hs
+++ b/compiler/GHC/Data/Graph/Ops.hs
@@ -218,8 +218,8 @@ addConflicts conflicts getClass
 
 addConflictSet1 :: Uniquable k
                 => k -> (k -> cls) -> UniqSet k
-                -> UniqFM (Node k cls color)
-                -> UniqFM (Node k cls color)
+                -> UniqFM k (Node k cls color)
+                -> UniqFM k (Node k cls color)
 addConflictSet1 u getClass set
  = case delOneFromUniqSet set u of
     set' -> adjustWithDefaultUFM
@@ -645,7 +645,7 @@ checkNode graph node
 
 slurpNodeConflictCount
         :: Graph k cls color
-        -> UniqFM (Int, Int)    -- ^ (conflict neighbours, num nodes with that many conflicts)
+        -> UniqFM Int (Int, Int)    -- ^ (conflict neighbours, num nodes with that many conflicts)
 
 slurpNodeConflictCount graph
         = addListToUFM_C
@@ -676,7 +676,7 @@ setColor u color
 adjustWithDefaultUFM
         :: Uniquable k
         => (a -> a) -> a -> k
-        -> UniqFM a -> UniqFM a
+        -> UniqFM k a -> UniqFM k a
 
 adjustWithDefaultUFM f def k map
         = addToUFM_C
@@ -689,7 +689,7 @@ adjustWithDefaultUFM f def k map
 adjustUFM_C
         :: Uniquable k
         => (a -> a)
-        -> k -> UniqFM a -> UniqFM a
+        -> k -> UniqFM k a -> UniqFM k a
 
 adjustUFM_C f k map
  = case lookupUFM map k of