Perturbate Unique keys in UniqFMwip/mix-uniques

3 files changed, 82 insertions, 45 deletions

diff --git a/compiler/GHC/Types/Unique/DFM.hs b/compiler/GHC/Types/Unique/DFM.hs
index f3009de2a2..337ecfa02f 100644
--- a/compiler/GHC/Types/Unique/DFM.hs
+++ b/compiler/GHC/Types/Unique/DFM.hs
@@ -69,7 +69,7 @@ module GHC.Types.Unique.DFM (
 
 import GHC.Prelude
 
-import GHC.Types.Unique ( Uniquable(..), Unique, getKey )
+import GHC.Types.Unique ( Uniquable(..), Unique )
 import GHC.Utils.Outputable
 
 import qualified Data.IntMap as M
@@ -78,7 +78,7 @@ import Data.Functor.Classes (Eq1 (..))
 import Data.List (sortBy)
 import Data.Function (on)
 import qualified Data.Semigroup as Semi
-import GHC.Types.Unique.FM (UniqFM, nonDetUFMToList, ufmToIntMap, unsafeIntMapToUFM)
+import GHC.Types.Unique.FM (UniqFM, getMixedKey, getUnmixedUnique, nonDetUFMToList, ufmToIntMap, unsafeIntMapToUFM)
 import Unsafe.Coerce
 
 -- Note [Deterministic UniqFM]
@@ -166,7 +166,7 @@ emptyUDFM :: UniqDFM key elt
 emptyUDFM = UDFM M.empty 0
 
 unitUDFM :: Uniquable key => key -> elt -> UniqDFM key elt
-unitUDFM k v = UDFM (M.singleton (getKey $ getUnique k) (TaggedVal v 0)) 1
+unitUDFM k v = UDFM (M.singleton (getMixedKey $ getUnique k) (TaggedVal v 0)) 1
 
 -- The new binding always goes to the right of existing ones
 addToUDFM :: Uniquable key => UniqDFM key elt -> key -> elt  -> UniqDFM key elt
@@ -175,7 +175,7 @@ addToUDFM m k v = addToUDFM_Directly m (getUnique k) v
 -- The new binding always goes to the right of existing ones
 addToUDFM_Directly :: UniqDFM key elt -> Unique -> elt -> UniqDFM key elt
 addToUDFM_Directly (UDFM m i) u v
-  = UDFM (M.insertWith tf (getKey u) (TaggedVal v i) m) (i + 1)
+  = UDFM (M.insertWith tf (getMixedKey u) (TaggedVal v i) m) (i + 1)
   where
     tf (TaggedVal new_v _) (TaggedVal _ old_i) = TaggedVal new_v old_i
       -- Keep the old tag, but insert the new value
@@ -188,7 +188,7 @@ addToUDFM_C_Directly
   -> Unique -> elt
   -> UniqDFM key elt
 addToUDFM_C_Directly f (UDFM m i) u v
-  = UDFM (M.insertWith tf (getKey u) (TaggedVal v i) m) (i + 1)
+  = UDFM (M.insertWith tf (getMixedKey u) (TaggedVal v i) m) (i + 1)
     where
       tf (TaggedVal new_v _) (TaggedVal old_v old_i)
          = TaggedVal (f old_v new_v) old_i
@@ -214,7 +214,7 @@ addListToUDFM_Directly_C
 addListToUDFM_Directly_C f = foldl' (\m (k, v) -> addToUDFM_C_Directly f m k v)
 
 delFromUDFM :: Uniquable key => UniqDFM key elt -> key -> UniqDFM key elt
-delFromUDFM (UDFM m i) k = UDFM (M.delete (getKey $ getUnique k) m) i
+delFromUDFM (UDFM m i) k = UDFM (M.delete (getMixedKey $ getUnique k) m) i
 
 plusUDFM_C :: (elt -> elt -> elt) -> UniqDFM key elt -> UniqDFM key elt -> UniqDFM key elt
 plusUDFM_C f udfml@(UDFM _ i) udfmr@(UDFM _ j)
@@ -272,13 +272,13 @@ insertUDFMIntoLeft_C f udfml udfmr =
   addListToUDFM_Directly_C f udfml $ udfmToList udfmr
 
 lookupUDFM :: Uniquable key => UniqDFM key elt -> key -> Maybe elt
-lookupUDFM (UDFM m _i) k = taggedFst `fmap` M.lookup (getKey $ getUnique k) m
+lookupUDFM (UDFM m _i) k = taggedFst `fmap` M.lookup (getMixedKey $ getUnique k) m
 
 lookupUDFM_Directly :: UniqDFM key elt -> Unique -> Maybe elt
-lookupUDFM_Directly (UDFM m _i) k = taggedFst `fmap` M.lookup (getKey k) m
+lookupUDFM_Directly (UDFM m _i) k = taggedFst `fmap` M.lookup (getMixedKey k) m
 
 elemUDFM :: Uniquable key => key -> UniqDFM key elt -> Bool
-elemUDFM k (UDFM m _i) = M.member (getKey $ getUnique k) m
+elemUDFM k (UDFM m _i) = M.member (getMixedKey $ getUnique k) m
 
 -- | Performs a deterministic fold over the UniqDFM.
 -- It's O(n log n) while the corresponding function on `UniqFM` is O(n).
@@ -304,13 +304,13 @@ filterUDFM p (UDFM m i) = UDFM (M.filter (\(TaggedVal v _) -> p v) m) i
 filterUDFM_Directly :: (Unique -> elt -> Bool) -> UniqDFM key elt -> UniqDFM key elt
 filterUDFM_Directly p (UDFM m i) = UDFM (M.filterWithKey p' m) i
   where
-  p' k (TaggedVal v _) = p (getUnique k) v
+  p' k (TaggedVal v _) = p (getUnmixedUnique k) v
 
 -- | Converts `UniqDFM` to a list, with elements in deterministic order.
 -- It's O(n log n) while the corresponding function on `UniqFM` is O(n).
 udfmToList :: UniqDFM key elt -> [(Unique, elt)]
 udfmToList (UDFM m _i) =
-  [ (getUnique k, taggedFst v)
+  [ (getUnmixedUnique k, taggedFst v)
   | (k, v) <- sortBy (compare `on` (taggedSnd . snd)) $ M.toList m ]
 
 -- Determines whether two 'UniqDFM's contain the same keys.
@@ -374,11 +374,11 @@ listToUDFM_Directly = foldl' (\m (u, v) -> addToUDFM_Directly m u v) emptyUDFM
 
 -- | Apply a function to a particular element
 adjustUDFM :: Uniquable key => (elt -> elt) -> UniqDFM key elt -> key -> UniqDFM key elt
-adjustUDFM f (UDFM m i) k = UDFM (M.adjust (fmap f) (getKey $ getUnique k) m) i
+adjustUDFM f (UDFM m i) k = UDFM (M.adjust (fmap f) (getMixedKey $ getUnique k) m) i
 
 -- | Apply a function to a particular element
 adjustUDFM_Directly :: (elt -> elt) -> UniqDFM key elt -> Unique -> UniqDFM key elt
-adjustUDFM_Directly f (UDFM m i) k = UDFM (M.adjust (fmap f) (getKey k) m) i
+adjustUDFM_Directly f (UDFM m i) k = UDFM (M.adjust (fmap f) (getMixedKey k) m) i
 
 -- | The expression (alterUDFM f k map) alters value x at k, or absence
 -- thereof. alterUDFM can be used to insert, delete, or update a value in
@@ -391,7 +391,7 @@ alterUDFM
   -> key                       -- new
   -> UniqDFM key elt               -- result
 alterUDFM f (UDFM m i) k =
-  UDFM (M.alter alterf (getKey $ getUnique k) m) (i + 1)
+  UDFM (M.alter alterf (getMixedKey $ getUnique k) m) (i + 1)
   where
   alterf Nothing = inject $ f Nothing
   alterf (Just (TaggedVal v _)) = inject $ f (Just v)
diff --git a/compiler/GHC/Types/Unique/FM.hs b/compiler/GHC/Types/Unique/FM.hs
index 6d13436169..2c2463d4f2 100644
--- a/compiler/GHC/Types/Unique/FM.hs
+++ b/compiler/GHC/Types/Unique/FM.hs
@@ -21,6 +21,7 @@ of arguments of combining function.
 -}
 
 {-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE RoleAnnotations #-}
 {-# LANGUAGE GeneralizedNewtypeDeriving #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE CPP #-}
@@ -30,6 +31,8 @@ module GHC.Types.Unique.FM (
         -- * Unique-keyed mappings
         UniqFM,           -- abstract type
         NonDetUniqFM(..), -- wrapper for opting into nondeterminism
+        getMixedKey,
+        getUnmixedUnique,
 
         -- ** Manipulating those mappings
         emptyUFM,
@@ -90,6 +93,7 @@ import qualified Data.IntSet as S
 import Data.Data
 import qualified Data.Semigroup as Semi
 import Data.Functor.Classes (Eq1 (..))
+import Data.Bits
 
 -- | A finite map from @uniques@ of one type to
 -- elements in another type.
@@ -104,6 +108,39 @@ newtype UniqFM key ele = UFM (M.IntMap ele)
   -- Nondeterministic Foldable and Traversable instances are accessible through
   -- use of the 'NonDetUniqFM' wrapper.
   -- See Note [Deterministic UniqFM] in GHC.Types.Unique.DFM to learn about determinism.
+type role UniqFM representational representational -- Don't allow coerces over the key
+
+-- | https://gist.github.com/degski/6e2069d6035ae04d5d6f64981c995ec2
+mix :: Word -> Int -> Int
+{-# INLINE mix #-}
+mix k x = fromIntegral $ f $ g $ f $ g $ f $ fromIntegral x
+  where
+    f y = (y `shiftR` s) `xor` y
+    g z = z * k
+    s = finiteBitSize k `shiftR` 1 -- 32 for 64 bit, 16 for 32 bit
+
+kFORWARD, kBACKWARD :: Word
+-- These are like "encryption" and "decryption" keys to mix
+#if UNIQUE_TAG_BITS == 8
+kFORWARD  = 0xD6E8FEB86659FD93
+kBACKWARD = 0xCFEE444D8B59A89B
+#else
+kFORWARD  = 0x45D9F3B
+kBACKWARD = 0x119DE1F3
+#endif
+enc, dec :: Int -> Int
+enc = mix kFORWARD
+dec = mix kBACKWARD
+{-# INLINE enc #-}
+{-# INLINE dec #-}
+
+getMixedKey :: Unique -> Int
+{-# INLINE getMixedKey #-}
+getMixedKey = enc . getKey
+
+getUnmixedUnique :: Int -> Unique
+{-# INLINE getUnmixedUnique #-}
+getUnmixedUnique = getUnique . dec
 
 emptyUFM :: UniqFM key elt
 emptyUFM = UFM M.empty
@@ -112,11 +149,11 @@ isNullUFM :: UniqFM key elt -> Bool
 isNullUFM (UFM m) = M.null m
 
 unitUFM :: Uniquable key => key -> elt -> UniqFM key elt
-unitUFM k v = UFM (M.singleton (getKey $ getUnique k) v)
+unitUFM k v = UFM (M.singleton (getMixedKey $ getUnique k) v)
 
 -- when you've got the Unique already
 unitDirectlyUFM :: Unique -> elt -> UniqFM key elt
-unitDirectlyUFM u v = UFM (M.singleton (getKey u) v)
+unitDirectlyUFM u v = UFM (M.singleton (getMixedKey u) v)
 
 -- zipToUFM ks vs = listToUFM (zip ks vs)
 -- This function exists because it's a common case (#18535), and
@@ -148,7 +185,7 @@ listToUFM_C
 listToUFM_C f = foldl' (\m (k, v) -> addToUFM_C f m k v) emptyUFM
 
 addToUFM :: Uniquable key => UniqFM key elt -> key -> elt  -> UniqFM key elt
-addToUFM (UFM m) k v = UFM (M.insert (getKey $ getUnique k) v m)
+addToUFM (UFM m) k v = UFM (M.insert (getMixedKey $ getUnique k) v m)
 
 addListToUFM :: Uniquable key => UniqFM key elt -> [(key,elt)] -> UniqFM key elt
 addListToUFM = foldl' (\m (k, v) -> addToUFM m k v)
@@ -157,7 +194,7 @@ addListToUFM_Directly :: UniqFM key elt -> [(Unique,elt)] -> UniqFM key elt
 addListToUFM_Directly = foldl' (\m (k, v) -> addToUFM_Directly m k v)
 
 addToUFM_Directly :: UniqFM key elt -> Unique -> elt -> UniqFM key elt
-addToUFM_Directly (UFM m) u v = UFM (M.insert (getKey u) v m)
+addToUFM_Directly (UFM m) u v = UFM (M.insert (getMixedKey u) v m)
 
 addToUFM_C
   :: Uniquable key
@@ -167,7 +204,7 @@ addToUFM_C
   -> UniqFM key elt           -- result
 -- Arguments of combining function of M.insertWith and addToUFM_C are flipped.
 addToUFM_C f (UFM m) k v =
-  UFM (M.insertWith (flip f) (getKey $ getUnique k) v m)
+  UFM (M.insertWith (flip f) (getMixedKey $ getUnique k) v m)
 
 addToUFM_Acc
   :: Uniquable key
@@ -177,7 +214,7 @@ addToUFM_Acc
   -> key -> elt             -- new
   -> UniqFM key elts            -- result
 addToUFM_Acc exi new (UFM m) k v =
-  UFM (M.insertWith (\_new old -> exi v old) (getKey $ getUnique k) (new v) m)
+  UFM (M.insertWith (\_new old -> exi v old) (getMixedKey $ getUnique k) (new v) m)
 
 alterUFM
   :: Uniquable key
@@ -185,7 +222,7 @@ alterUFM
   -> UniqFM key elt                -- old
   -> key                       -- new
   -> UniqFM key elt                -- result
-alterUFM f (UFM m) k = UFM (M.alter f (getKey $ getUnique k) m)
+alterUFM f (UFM m) k = UFM (M.alter f (getMixedKey $ getUnique k) m)
 
 -- | Add elements to the map, combining existing values with inserted ones using
 -- the given function.
@@ -197,13 +234,13 @@ addListToUFM_C
 addListToUFM_C f = foldl' (\m (k, v) -> addToUFM_C f m k v)
 
 adjustUFM :: Uniquable key => (elt -> elt) -> UniqFM key elt -> key -> UniqFM key elt
-adjustUFM f (UFM m) k = UFM (M.adjust f (getKey $ getUnique k) m)
+adjustUFM f (UFM m) k = UFM (M.adjust f (getMixedKey $ getUnique k) m)
 
 adjustUFM_Directly :: (elt -> elt) -> UniqFM key elt -> Unique -> UniqFM key elt
-adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getKey u) m)
+adjustUFM_Directly f (UFM m) u = UFM (M.adjust f (getMixedKey u) m)
 
 delFromUFM :: Uniquable key => UniqFM key elt -> key    -> UniqFM key elt
-delFromUFM (UFM m) k = UFM (M.delete (getKey $ getUnique k) m)
+delFromUFM (UFM m) k = UFM (M.delete (getMixedKey $ getUnique k) m)
 
 delListFromUFM :: Uniquable key => UniqFM key elt -> [key] -> UniqFM key elt
 delListFromUFM = foldl' delFromUFM
@@ -212,7 +249,7 @@ delListFromUFM_Directly :: UniqFM key elt -> [Unique] -> UniqFM key elt
 delListFromUFM_Directly = foldl' delFromUFM_Directly
 
 delFromUFM_Directly :: UniqFM key elt -> Unique -> UniqFM key elt
-delFromUFM_Directly (UFM m) u = UFM (M.delete (getKey u) m)
+delFromUFM_Directly (UFM m) u = UFM (M.delete (getMixedKey u) m)
 
 -- Bindings in right argument shadow those in the left
 plusUFM :: UniqFM key elt -> UniqFM key elt -> UniqFM key elt
@@ -303,13 +340,13 @@ mapUFM :: (elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2
 mapUFM f (UFM m) = UFM (M.map f m)
 
 mapUFM_Directly :: (Unique -> elt1 -> elt2) -> UniqFM key elt1 -> UniqFM key elt2
-mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnique) m)
+mapUFM_Directly f (UFM m) = UFM (M.mapWithKey (f . getUnmixedUnique) m)
 
 filterUFM :: (elt -> Bool) -> UniqFM key elt -> UniqFM key elt
 filterUFM p (UFM m) = UFM (M.filter p m)
 
 filterUFM_Directly :: (Unique -> elt -> Bool) -> UniqFM key elt -> UniqFM key elt
-filterUFM_Directly p (UFM m) = UFM (M.filterWithKey (p . getUnique) m)
+filterUFM_Directly p (UFM m) = UFM (M.filterWithKey (p . getUnmixedUnique) m)
 
 partitionUFM :: (elt -> Bool) -> UniqFM key elt -> (UniqFM key elt, UniqFM key elt)
 partitionUFM p (UFM m) =
@@ -320,29 +357,29 @@ sizeUFM :: UniqFM key elt -> Int
 sizeUFM (UFM m) = M.size m
 
 elemUFM :: Uniquable key => key -> UniqFM key elt -> Bool
-elemUFM k (UFM m) = M.member (getKey $ getUnique k) m
+elemUFM k (UFM m) = M.member (getMixedKey $ getUnique k) m
 
 elemUFM_Directly :: Unique -> UniqFM key elt -> Bool
-elemUFM_Directly u (UFM m) = M.member (getKey u) m
+elemUFM_Directly u (UFM m) = M.member (getMixedKey u) m
 
 lookupUFM :: Uniquable key => UniqFM key elt -> key -> Maybe elt
-lookupUFM (UFM m) k = M.lookup (getKey $ getUnique k) m
+lookupUFM (UFM m) k = M.lookup (getMixedKey $ getUnique k) m
 
 -- when you've got the Unique already
 lookupUFM_Directly :: UniqFM key elt -> Unique -> Maybe elt
-lookupUFM_Directly (UFM m) u = M.lookup (getKey u) m
+lookupUFM_Directly (UFM m) u = M.lookup (getMixedKey u) m
 
 lookupWithDefaultUFM :: Uniquable key => UniqFM key elt -> elt -> key -> elt
-lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getKey $ getUnique k) m
+lookupWithDefaultUFM (UFM m) v k = M.findWithDefault v (getMixedKey $ getUnique k) m
 
 lookupWithDefaultUFM_Directly :: UniqFM key elt -> elt -> Unique -> elt
-lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getKey u) m
+lookupWithDefaultUFM_Directly (UFM m) v u = M.findWithDefault v (getMixedKey u) m
 
 eltsUFM :: UniqFM key elt -> [elt]
 eltsUFM (UFM m) = M.elems m
 
 ufmToSet_Directly :: UniqFM key elt -> S.IntSet
-ufmToSet_Directly (UFM m) = M.keysSet m
+ufmToSet_Directly m = S.fromList $ map getKey $ nonDetKeysUFM m
 
 anyUFM :: (elt -> Bool) -> UniqFM key elt -> Bool
 anyUFM p (UFM m) = M.foldr ((||) . p) False m
@@ -366,7 +403,7 @@ nonDetEltsUFM (UFM m) = M.elems m
 -- If you use this please provide a justification why it doesn't introduce
 -- nondeterminism.
 nonDetKeysUFM :: UniqFM key elt -> [Unique]
-nonDetKeysUFM (UFM m) = map getUnique $ M.keys m
+nonDetKeysUFM (UFM m) = map getUnmixedUnique $ M.keys m
 
 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
 -- If you use this please provide a justification why it doesn't introduce
@@ -378,13 +415,13 @@ nonDetStrictFoldUFM k z (UFM m) = M.foldl' (flip k) z m
 -- If you use this please provide a justification why it doesn't introduce
 -- nondeterminism.
 nonDetStrictFoldUFM_Directly:: (Unique -> elt -> a -> a) -> a -> UniqFM key elt -> a
-nonDetStrictFoldUFM_Directly k z (UFM m) = M.foldlWithKey' (\z' i x -> k (getUnique i) x z') z m
+nonDetStrictFoldUFM_Directly k z (UFM m) = M.foldlWithKey' (\z' i x -> k (getUnmixedUnique i) x z') z m
 
 -- See Note [Deterministic UniqFM] to learn about nondeterminism.
 -- If you use this please provide a justification why it doesn't introduce
 -- nondeterminism.
 nonDetUFMToList :: UniqFM key elt -> [(Unique, elt)]
-nonDetUFMToList (UFM m) = map (\(k, v) -> (getUnique k, v)) $ M.toList m
+nonDetUFMToList (UFM m) = map (\(k, v) -> (getUnmixedUnique k, v)) $ M.toList m
 
 -- | A wrapper around 'UniqFM' with the sole purpose of informing call sites
 -- that the provided 'Foldable' and 'Traversable' instances are
diff --git a/compiler/GHC/Types/Var/Env.hs b/compiler/GHC/Types/Var/Env.hs
index ed58c413f4..583deec316 100644
--- a/compiler/GHC/Types/Var/Env.hs
+++ b/compiler/GHC/Types/Var/Env.hs
@@ -200,14 +200,14 @@ uniqAway' in_scope var
 -- given 'InScopeSet'. This must be used very carefully since one can very easily
 -- introduce non-unique 'Unique's this way. See Note [Local uniques].
 unsafeGetFreshLocalUnique :: InScopeSet -> Unique
-unsafeGetFreshLocalUnique (InScope set)
-  | Just (uniq,_) <- IntMap.lookupLT (getKey maxLocalUnique) (ufmToIntMap $ getUniqSet set)
-  , let uniq' = mkLocalUnique uniq
-  , not $ uniq' `ltUnique` minLocalUnique
-  = incrUnique uniq'
-
-  | otherwise
-  = minLocalUnique
+unsafeGetFreshLocalUnique (InScope set) = go (getMixedKey (getUnique (sizeUniqSet set))) -- much stuff to fix here
+  where
+    go n
+      | let uniq = mkLocalUnique n
+      , Nothing <- IntMap.lookup (getMixedKey $ uniq) (ufmToIntMap $ getUniqSet set)
+      = uniq
+      | otherwise
+      = go (getMixedKey $ getUnique (n+1)) -- hopefully this will eventually traverse the whole space
 
 {-
 ************************************************************************