UniqMap implementation.

Signed-off-by: Edward Z. Yang <ezyang@cs.stanford.edu>

2 files changed, 211 insertions, 0 deletions

diff --git a/compiler/ghc.cabal.in b/compiler/ghc.cabal.in
index 8c9bc3bfaa..fc8dcd98bf 100644
--- a/compiler/ghc.cabal.in
+++ b/compiler/ghc.cabal.in
@@ -508,6 +508,7 @@ Library
         UniqDFM
         UniqDSet
         UniqFM
+        UniqMap
         UniqSet
         Util
         Vectorise.Builtins.Base
diff --git a/compiler/utils/UniqMap.hs b/compiler/utils/UniqMap.hs
new file mode 100644
index 0000000000..012409b8c2
--- /dev/null
+++ b/compiler/utils/UniqMap.hs
@@ -0,0 +1,210 @@
+{-# LANGUAGE RoleAnnotations #-}
+{-# LANGUAGE TupleSections #-}
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE DeriveFunctor #-}
+{-# OPTIONS_GHC -Wall #-}
+
+-- Like 'UniqFM', these are maps for keys which are Uniquable.
+-- Unlike 'UniqFM', these maps also remember their keys, which
+-- makes them a much better drop in replacement for 'Data.Map.Map'.
+--
+-- Key preservation is right-biased.
+module UniqMap (
+    UniqMap,
+    emptyUniqMap,
+    isNullUniqMap,
+    unitUniqMap,
+    listToUniqMap,
+    listToUniqMap_C,
+    addToUniqMap,
+    addListToUniqMap,
+    addToUniqMap_C,
+    addToUniqMap_Acc,
+    alterUniqMap,
+    addListToUniqMap_C,
+    adjustUniqMap,
+    delFromUniqMap,
+    delListFromUniqMap,
+    plusUniqMap,
+    plusUniqMap_C,
+    plusMaybeUniqMap_C,
+    plusUniqMapList,
+    minusUniqMap,
+    intersectUniqMap,
+    disjointUniqMap,
+    mapUniqMap,
+    filterUniqMap,
+    partitionUniqMap,
+    sizeUniqMap,
+    elemUniqMap,
+    lookupUniqMap,
+    lookupWithDefaultUniqMap,
+    anyUniqMap,
+    allUniqMap,
+    -- Non-deterministic functions omitted
+) where
+
+import UniqFM
+
+import Unique
+import Outputable
+
+#if __GLASGOW_HASKELL__ > 710
+import Data.Semigroup   ( Semigroup(..) )
+#endif
+import Data.Coerce
+import Data.Maybe
+import Data.Typeable
+import Data.Data
+
+-- | Maps indexed by 'Uniquable' keys
+newtype UniqMap k a = UniqMap (UniqFM (k, a))
+    deriving (Data, Eq, Functor, Typeable)
+type role UniqMap nominal representational
+
+#if __GLASGOW_HASKELL__ > 710
+instance Semigroup (UniqMap k a) where
+  (<>) = plusUniqMap
+#endif
+
+instance Monoid (UniqMap k a) where
+    mempty = emptyUniqMap
+    mappend = plusUniqMap
+
+instance (Outputable k, Outputable a) => Outputable (UniqMap k a) where
+    ppr (UniqMap m) =
+        brackets $ fsep $ punctuate comma $
+        [ ppr k <+> text "->" <+> ppr v
+        | (k, v) <- eltsUFM m ]
+
+liftC :: (a -> a -> a) -> (k, a) -> (k, a) -> (k, a)
+liftC f (_, v) (k', v') = (k', f v v')
+
+emptyUniqMap :: UniqMap k a
+emptyUniqMap = UniqMap emptyUFM
+
+isNullUniqMap :: UniqMap k a -> Bool
+isNullUniqMap (UniqMap m) = isNullUFM m
+
+unitUniqMap :: Uniquable k => k -> a -> UniqMap k a
+unitUniqMap k v = UniqMap (unitUFM k (k, v))
+
+listToUniqMap :: Uniquable k => [(k,a)] -> UniqMap k a
+listToUniqMap kvs = UniqMap (listToUFM [ (k,(k,v)) | (k,v) <- kvs])
+
+listToUniqMap_C :: Uniquable k => (a -> a -> a) -> [(k,a)] -> UniqMap k a
+listToUniqMap_C f kvs = UniqMap $
+    listToUFM_C (liftC f) [ (k,(k,v)) | (k,v) <- kvs]
+
+addToUniqMap :: Uniquable k => UniqMap k a -> k -> a -> UniqMap k a
+addToUniqMap (UniqMap m) k v = UniqMap $ addToUFM m k (k, v)
+
+addListToUniqMap :: Uniquable k => UniqMap k a -> [(k,a)] -> UniqMap k a
+addListToUniqMap (UniqMap m) kvs = UniqMap $
+    addListToUFM m [(k,(k,v)) | (k,v) <- kvs]
+
+addToUniqMap_C :: Uniquable k
+               => (a -> a -> a)
+               -> UniqMap k a
+               -> k
+               -> a
+               -> UniqMap k a
+addToUniqMap_C f (UniqMap m) k v = UniqMap $
+    addToUFM_C (liftC f) m k (k, v)
+
+addToUniqMap_Acc :: Uniquable k
+                 => (b -> a -> a)
+                 -> (b -> a)
+                 -> UniqMap k a
+                 -> k
+                 -> b
+                 -> UniqMap k a
+addToUniqMap_Acc exi new (UniqMap m) k0 v0 = UniqMap $
+    addToUFM_Acc (\b (k, v) -> (k, exi b v))
+                 (\b -> (k0, new b))
+                 m k0 v0
+
+alterUniqMap :: Uniquable k
+             => (Maybe a -> Maybe a)
+             -> UniqMap k a
+             -> k
+             -> UniqMap k a
+alterUniqMap f (UniqMap m) k = UniqMap $
+    alterUFM (fmap (k,) . f . fmap snd) m k
+
+addListToUniqMap_C
+    :: Uniquable k
+    => (a -> a -> a)
+    -> UniqMap k a
+    -> [(k, a)]
+    -> UniqMap k a
+addListToUniqMap_C f (UniqMap m) kvs = UniqMap $
+    addListToUFM_C (liftC f) m
+        [(k,(k,v)) | (k,v) <- kvs]
+
+adjustUniqMap
+    :: Uniquable k
+    => (a -> a)
+    -> UniqMap k a
+    -> k
+    -> UniqMap k a
+adjustUniqMap f (UniqMap m) k = UniqMap $
+    adjustUFM (\(_,v) -> (k,f v)) m k
+
+delFromUniqMap :: Uniquable k => UniqMap k a -> k -> UniqMap k a
+delFromUniqMap (UniqMap m) k = UniqMap $ delFromUFM m k
+
+delListFromUniqMap :: Uniquable k => UniqMap k a -> [k] -> UniqMap k a
+delListFromUniqMap (UniqMap m) ks = UniqMap $ delListFromUFM m ks
+
+plusUniqMap :: UniqMap k a -> UniqMap k a -> UniqMap k a
+plusUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ plusUFM m1 m2
+
+plusUniqMap_C :: (a -> a -> a) -> UniqMap k a -> UniqMap k a -> UniqMap k a
+plusUniqMap_C f (UniqMap m1) (UniqMap m2) = UniqMap $
+    plusUFM_C (liftC f) m1 m2
+
+plusMaybeUniqMap_C :: (a -> a -> Maybe a) -> UniqMap k a -> UniqMap k a -> UniqMap k a
+plusMaybeUniqMap_C f (UniqMap m1) (UniqMap m2) = UniqMap $
+    plusMaybeUFM_C (\(_, v) (k', v') -> fmap (k',) (f v v')) m1 m2
+
+plusUniqMapList :: [UniqMap k a] -> UniqMap k a
+plusUniqMapList xs = UniqMap $ plusUFMList (coerce xs)
+
+minusUniqMap :: UniqMap k a -> UniqMap k b -> UniqMap k a
+minusUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ minusUFM m1 m2
+
+intersectUniqMap :: UniqMap k a -> UniqMap k b -> UniqMap k a
+intersectUniqMap (UniqMap m1) (UniqMap m2) = UniqMap $ intersectUFM m1 m2
+
+disjointUniqMap :: UniqMap k a -> UniqMap k b -> Bool
+disjointUniqMap (UniqMap m1) (UniqMap m2) = disjointUFM m1 m2
+
+mapUniqMap :: (a -> b) -> UniqMap k a -> UniqMap k b
+mapUniqMap f (UniqMap m) = UniqMap $ mapUFM (fmap f) m -- (,) k instance
+
+filterUniqMap :: (a -> Bool) -> UniqMap k a -> UniqMap k a
+filterUniqMap f (UniqMap m) = UniqMap $ filterUFM (f . snd) m
+
+partitionUniqMap :: (a -> Bool) -> UniqMap k a -> (UniqMap k a, UniqMap k a)
+partitionUniqMap f (UniqMap m) =
+    coerce $ partitionUFM (f . snd) m
+
+sizeUniqMap :: UniqMap k a -> Int
+sizeUniqMap (UniqMap m) = sizeUFM m
+
+elemUniqMap :: Uniquable k => k -> UniqMap k a -> Bool
+elemUniqMap k (UniqMap m) = elemUFM k m
+
+lookupUniqMap :: Uniquable k => UniqMap k a -> k -> Maybe a
+lookupUniqMap (UniqMap m) k = fmap snd (lookupUFM m k)
+
+lookupWithDefaultUniqMap :: Uniquable k => UniqMap k a -> a -> k -> a
+lookupWithDefaultUniqMap (UniqMap m) a k = fromMaybe a (fmap snd (lookupUFM m k))
+
+anyUniqMap :: (a -> Bool) -> UniqMap k a -> Bool
+anyUniqMap f (UniqMap m) = anyUFM (f . snd) m
+
+allUniqMap :: (a -> Bool) -> UniqMap k a -> Bool
+allUniqMap f (UniqMap m) = allUFM (f . snd) m