Make `Functor` a superclass of `TrieMap`, which lets us derive the `map` functions.

2 files changed, 77 insertions, 67 deletions

diff --git a/compiler/GHC/Core/Map/Expr.hs b/compiler/GHC/Core/Map/Expr.hs
index 60ee2c94b5..61fca8353a 100644
--- a/compiler/GHC/Core/Map/Expr.hs
+++ b/compiler/GHC/Core/Map/Expr.hs
@@ -109,13 +109,17 @@ See also Note [Empty case alternatives] in GHC.Core.
 -- is the type you want.
 newtype CoreMap a = CoreMap (CoreMapG a)
 
+-- TODO(22292): derive
+instance Functor CoreMap where
+    fmap f = \ (CoreMap m) -> CoreMap (fmap f m)
+    {-# INLINE fmap #-}
+
 instance TrieMap CoreMap where
     type Key CoreMap = CoreExpr
     emptyTM = CoreMap emptyTM
     lookupTM k (CoreMap m) = lookupTM (deBruijnize k) m
     alterTM k f (CoreMap m) = CoreMap (alterTM (deBruijnize k) f m)
     foldTM k (CoreMap m) = foldTM k m
-    mapTM f (CoreMap m) = CoreMap (mapTM f m)
     filterTM f (CoreMap m) = CoreMap (filterTM f m)
 
 -- | @CoreMapG a@ is a map from @DeBruijn CoreExpr@ to @a@.  The extended
@@ -248,30 +252,27 @@ emptyE = CM { cm_var = emptyTM, cm_lit = emptyTM
             , cm_letr = emptyTM, cm_case = emptyTM
             , cm_ecase = emptyTM, cm_tick = emptyTM }
 
+-- TODO(22292): derive
+instance Functor CoreMapX where
+    fmap f CM
+      { cm_var = cvar, cm_lit = clit, cm_co = cco, cm_type = ctype, cm_cast = ccast
+      , cm_app = capp, cm_lam = clam, cm_letn = cletn, cm_letr = cletr, cm_case = ccase
+      , cm_ecase = cecase, cm_tick = ctick } = CM
+      { cm_var = fmap f cvar, cm_lit = fmap f clit, cm_co = fmap f cco, cm_type = fmap f ctype
+      , cm_cast = fmap (fmap f) ccast, cm_app = fmap (fmap f) capp, cm_lam = fmap (fmap f) clam
+      , cm_letn = fmap (fmap (fmap f)) cletn, cm_letr = fmap (fmap (fmap f)) cletr
+      , cm_case = fmap (fmap f) ccase, cm_ecase = fmap (fmap f) cecase
+      , cm_tick = fmap (fmap f) ctick }
+
 instance TrieMap CoreMapX where
    type Key CoreMapX = DeBruijn CoreExpr
    emptyTM  = emptyE
    lookupTM = lkE
    alterTM  = xtE
    foldTM   = fdE
-   mapTM    = mapE
    filterTM = ftE
 
 --------------------------
-mapE :: (a->b) -> CoreMapX a -> CoreMapX b
-mapE f (CM { cm_var = cvar, cm_lit = clit
-           , cm_co = cco, cm_type = ctype
-           , cm_cast = ccast , cm_app = capp
-           , cm_lam = clam, cm_letn = cletn
-           , cm_letr = cletr, cm_case = ccase
-           , cm_ecase = cecase, cm_tick = ctick })
-  = CM { cm_var = mapTM f cvar, cm_lit = mapTM f clit
-       , cm_co = mapTM f cco, cm_type = mapTM f ctype
-       , cm_cast = mapTM (mapTM f) ccast, cm_app = mapTM (mapTM f) capp
-       , cm_lam = mapTM (mapTM f) clam, cm_letn = mapTM (mapTM (mapTM f)) cletn
-       , cm_letr = mapTM (mapTM (mapTM f)) cletr, cm_case = mapTM (mapTM f) ccase
-       , cm_ecase = mapTM (mapTM f) cecase, cm_tick = mapTM (mapTM f) ctick }
-
 ftE :: (a->Bool) -> CoreMapX a -> CoreMapX a
 ftE f (CM { cm_var = cvar, cm_lit = clit
           , cm_co = cco, cm_type = ctype
@@ -281,10 +282,10 @@ ftE f (CM { cm_var = cvar, cm_lit = clit
           , cm_ecase = cecase, cm_tick = ctick })
   = CM { cm_var = filterTM f cvar, cm_lit = filterTM f clit
        , cm_co = filterTM f cco, cm_type = filterTM f ctype
-       , cm_cast = mapTM (filterTM f) ccast, cm_app = mapTM (filterTM f) capp
-       , cm_lam = mapTM (filterTM f) clam, cm_letn = mapTM (mapTM (filterTM f)) cletn
-       , cm_letr = mapTM (mapTM (filterTM f)) cletr, cm_case = mapTM (filterTM f) ccase
-       , cm_ecase = mapTM (filterTM f) cecase, cm_tick = mapTM (filterTM f) ctick }
+       , cm_cast = fmap (filterTM f) ccast, cm_app = fmap (filterTM f) capp
+       , cm_lam = fmap (filterTM f) clam, cm_letn = fmap (fmap (filterTM f)) cletn
+       , cm_letr = fmap (fmap (filterTM f)) cletr, cm_case = fmap (filterTM f) ccase
+       , cm_ecase = fmap (filterTM f) cecase, cm_tick = fmap (filterTM f) ctick }
 
 --------------------------
 lookupCoreMap :: CoreMap a -> CoreExpr -> Maybe a
@@ -394,6 +395,11 @@ data AltMap a   -- A single alternative
        , am_data  :: DNameEnv (CoreMapG a)
        , am_lit   :: LiteralMap (CoreMapG a) }
 
+-- TODO(22292): derive
+instance Functor AltMap where
+    fmap f AM { am_deflt = adeflt, am_data = adata, am_lit = alit } = AM
+      { am_deflt = fmap f adeflt, am_data = fmap (fmap f) adata, am_lit = fmap (fmap f) alit }
+
 instance TrieMap AltMap where
    type Key AltMap = CoreAlt
    emptyTM  = AM { am_deflt = emptyTM
@@ -402,7 +408,6 @@ instance TrieMap AltMap where
    lookupTM = lkA emptyCME
    alterTM  = xtA emptyCME
    foldTM   = fdA
-   mapTM    = mapA
    filterTM = ftA
 
 instance Eq (DeBruijn CoreAlt) where
@@ -416,17 +421,11 @@ instance Eq (DeBruijn CoreAlt) where
           D (extendCMEs env1 bs1) rhs1 == D (extendCMEs env2 bs2) rhs2
     go _ _ = False
 
-mapA :: (a->b) -> AltMap a -> AltMap b
-mapA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
-  = AM { am_deflt = mapTM f adeflt
-       , am_data = mapTM (mapTM f) adata
-       , am_lit = mapTM (mapTM f) alit }
-
 ftA :: (a->Bool) -> AltMap a -> AltMap a
 ftA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
   = AM { am_deflt = filterTM f adeflt
-       , am_data = mapTM (filterTM f) adata
-       , am_lit = mapTM (filterTM f) alit }
+       , am_data = fmap (filterTM f) adata
+       , am_lit = fmap (filterTM f) alit }
 
 lkA :: CmEnv -> CoreAlt -> AltMap a -> Maybe a
 lkA env (Alt DEFAULT      _  rhs) = am_deflt >.> lkG (D env rhs)
diff --git a/compiler/GHC/Core/Map/Type.hs b/compiler/GHC/Core/Map/Type.hs
index 1617d93991..08d7fcf4e0 100644
--- a/compiler/GHC/Core/Map/Type.hs
+++ b/compiler/GHC/Core/Map/Type.hs
@@ -83,25 +83,33 @@ import GHC.Data.Maybe
 -- just look up the coercion's type.
 newtype CoercionMap a = CoercionMap (CoercionMapG a)
 
+-- TODO(22292): derive
+instance Functor CoercionMap where
+    fmap f = \ (CoercionMap m) -> CoercionMap (fmap f m)
+    {-# INLINE fmap #-}
+
 instance TrieMap CoercionMap where
    type Key CoercionMap = Coercion
    emptyTM                     = CoercionMap emptyTM
    lookupTM k  (CoercionMap m) = lookupTM (deBruijnize k) m
    alterTM k f (CoercionMap m) = CoercionMap (alterTM (deBruijnize k) f m)
    foldTM k    (CoercionMap m) = foldTM k m
-   mapTM f     (CoercionMap m) = CoercionMap (mapTM f m)
    filterTM f  (CoercionMap m) = CoercionMap (filterTM f m)
 
 type CoercionMapG = GenMap CoercionMapX
 newtype CoercionMapX a = CoercionMapX (TypeMapX a)
 
+-- TODO(22292): derive
+instance Functor CoercionMapX where
+    fmap f = \ (CoercionMapX core_tm) -> CoercionMapX (fmap f core_tm)
+    {-# INLINE fmap #-}
+
 instance TrieMap CoercionMapX where
   type Key CoercionMapX = DeBruijn Coercion
   emptyTM = CoercionMapX emptyTM
   lookupTM = lkC
   alterTM  = xtC
   foldTM f (CoercionMapX core_tm) = foldTM f core_tm
-  mapTM f (CoercionMapX core_tm)  = CoercionMapX (mapTM f core_tm)
   filterTM f (CoercionMapX core_tm) = CoercionMapX (filterTM f core_tm)
 
 instance Eq (DeBruijn Coercion) where
@@ -172,13 +180,21 @@ trieMapView ty
   | Just ty' <- tcView ty = Just ty'
 trieMapView _ = Nothing
 
+-- TODO(22292): derive
+instance Functor TypeMapX where
+    fmap f TM
+      { tm_var = tvar, tm_app = tapp, tm_tycon = ttycon, tm_funty = tfunty, tm_forall = tforall
+      , tm_tylit = tlit, tm_coerce = tcoerce } = TM
+      { tm_var = fmap f tvar, tm_app = fmap (fmap f) tapp, tm_tycon = fmap f ttycon
+      , tm_funty = fmap (fmap (fmap f)) tfunty, tm_forall = fmap (fmap f) tforall
+      , tm_tylit  = fmap f tlit, tm_coerce = fmap f tcoerce }
+
 instance TrieMap TypeMapX where
    type Key TypeMapX = DeBruijn Type
    emptyTM  = emptyT
    lookupTM = lkT
    alterTM  = xtT
    foldTM   = fdT
-   mapTM    = mapT
    filterTM = filterT
 
 instance Eq (DeBruijn Type) where
@@ -313,18 +329,6 @@ emptyT = TM { tm_var  = emptyTM
             , tm_tylit  = emptyTyLitMap
             , tm_coerce = Nothing }
 
-mapT :: (a->b) -> TypeMapX a -> TypeMapX b
-mapT f (TM { tm_var  = tvar, tm_app = tapp, tm_tycon = ttycon
-           , tm_funty = tfunty, tm_forall = tforall, tm_tylit = tlit
-           , tm_coerce = tcoerce })
-  = TM { tm_var    = mapTM f tvar
-       , tm_app    = mapTM (mapTM f) tapp
-       , tm_tycon  = mapTM f ttycon
-       , tm_funty  = mapTM (mapTM (mapTM f)) tfunty
-       , tm_forall = mapTM (mapTM f) tforall
-       , tm_tylit  = mapTM f tlit
-       , tm_coerce = fmap f tcoerce }
-
 -----------------
 lkT :: DeBruijn Type -> TypeMapX a -> Maybe a
 lkT (D env ty) m = go ty m
@@ -382,10 +386,10 @@ filterT f (TM { tm_var  = tvar, tm_app = tapp, tm_tycon = ttycon
               , tm_funty = tfunty, tm_forall = tforall, tm_tylit = tlit
               , tm_coerce = tcoerce })
   = TM { tm_var    = filterTM f tvar
-       , tm_app    = mapTM (filterTM f) tapp
+       , tm_app    = fmap (filterTM f) tapp
        , tm_tycon  = filterTM f ttycon
-       , tm_funty  = mapTM (mapTM (filterTM f)) tfunty
-       , tm_forall = mapTM (filterTM f) tforall
+       , tm_funty  = fmap (fmap (filterTM f)) tfunty
+       , tm_forall = fmap (filterTM f) tforall
        , tm_tylit  = filterTM f tlit
        , tm_coerce = filterMaybe f tcoerce }
 
@@ -395,22 +399,22 @@ data TyLitMap a = TLM { tlm_number :: Map.Map Integer a
                       , tlm_char   :: Map.Map Char a
                       }
 
+-- TODO(22292): derive
+instance Functor TyLitMap where
+    fmap f TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc } = TLM
+      { tlm_number = Map.map f tn, tlm_string = mapUFM f ts, tlm_char = Map.map f tc }
+
 instance TrieMap TyLitMap where
    type Key TyLitMap = TyLit
    emptyTM  = emptyTyLitMap
    lookupTM = lkTyLit
    alterTM  = xtTyLit
    foldTM   = foldTyLit
-   mapTM    = mapTyLit
    filterTM = filterTyLit
 
 emptyTyLitMap :: TyLitMap a
 emptyTyLitMap = TLM { tlm_number = Map.empty, tlm_string = emptyUFM, tlm_char = Map.empty }
 
-mapTyLit :: (a->b) -> TyLitMap a -> TyLitMap b
-mapTyLit f (TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc })
-  = TLM { tlm_number = Map.map f tn, tlm_string = mapUFM f ts, tlm_char = Map.map f tc }
-
 lkTyLit :: TyLit -> TyLitMap a -> Maybe a
 lkTyLit l =
   case l of
@@ -439,6 +443,11 @@ filterTyLit f (TLM { tlm_number = tn, tlm_string = ts, tlm_char = tc })
 -- is the type you want. The keys in this map may have different kinds.
 newtype TypeMap a = TypeMap (TypeMapG (TypeMapG a))
 
+-- TODO(22292): derive
+instance Functor TypeMap where
+    fmap f = \ (TypeMap m) -> TypeMap (fmap (fmap f) m)
+    {-# INLINE fmap #-}
+
 lkTT :: DeBruijn Type -> TypeMap a -> Maybe a
 lkTT (D env ty) (TypeMap m) = lkG (D env $ typeKind ty) m
                           >>= lkG (D env ty)
@@ -456,8 +465,7 @@ instance TrieMap TypeMap where
     lookupTM k m = lkTT (deBruijnize k) m
     alterTM k f m = xtTT (deBruijnize k) f m
     foldTM k (TypeMap m) = foldTM (foldTM k) m
-    mapTM f (TypeMap m) = TypeMap (mapTM (mapTM f) m)
-    filterTM f (TypeMap m) = TypeMap (mapTM (filterTM f) m)
+    filterTM f (TypeMap m) = TypeMap (fmap (filterTM f) m)
 
 foldTypeMap :: (a -> b -> b) -> b -> TypeMap a -> b
 foldTypeMap k z m = foldTM k m z
@@ -488,8 +496,12 @@ mkDeBruijnContext = extendCMEs emptyCME
 
 -- | A 'LooseTypeMap' doesn't do a kind-check. Thus, when lookup up (t |> g),
 -- you'll find entries inserted under (t), even if (g) is non-reflexive.
-newtype LooseTypeMap a
-  = LooseTypeMap (TypeMapG a)
+newtype LooseTypeMap a = LooseTypeMap (TypeMapG a)
+
+-- TODO(22292): derive
+instance Functor LooseTypeMap where
+    fmap f = \ (LooseTypeMap m) -> LooseTypeMap (fmap f m)
+    {-# INLINE fmap #-}
 
 instance TrieMap LooseTypeMap where
   type Key LooseTypeMap = Type
@@ -497,7 +509,6 @@ instance TrieMap LooseTypeMap where
   lookupTM k (LooseTypeMap m) = lookupTM (deBruijnize k) m
   alterTM k f (LooseTypeMap m) = LooseTypeMap (alterTM (deBruijnize k) f m)
   foldTM f (LooseTypeMap m) = foldTM f m
-  mapTM f (LooseTypeMap m) = LooseTypeMap (mapTM f m)
   filterTM f (LooseTypeMap m) = LooseTypeMap (filterTM f m)
 
 {-
@@ -566,18 +577,19 @@ instance Eq (DeBruijn a) => Eq (DeBruijn (Maybe a)) where
 -- of pairs are composition.
 data BndrMap a = BndrMap (TypeMapG (MaybeMap TypeMapG a))
 
+-- TODO(22292): derive
+instance Functor BndrMap where
+    fmap f = \ (BndrMap tm) -> BndrMap (fmap (fmap f) tm)
+    {-# INLINE fmap #-}
+
 instance TrieMap BndrMap where
    type Key BndrMap = Var
    emptyTM  = BndrMap emptyTM
    lookupTM = lkBndr emptyCME
    alterTM  = xtBndr emptyCME
    foldTM   = fdBndrMap
-   mapTM    = mapBndrMap
    filterTM = ftBndrMap
 
-mapBndrMap :: (a -> b) -> BndrMap a -> BndrMap b
-mapBndrMap f (BndrMap tm) = BndrMap (mapTM (mapTM f) tm)
-
 fdBndrMap :: (a -> b -> b) -> BndrMap a -> b -> b
 fdBndrMap f (BndrMap tm) = foldTM (foldTM f) tm
 
@@ -596,25 +608,24 @@ xtBndr env v xt (BndrMap tymap)  =
   BndrMap (tymap |> xtG (D env (varType v)) |>> (alterTM (D env <$> varMultMaybe v) xt))
 
 ftBndrMap :: (a -> Bool) -> BndrMap a -> BndrMap a
-ftBndrMap f (BndrMap tm) = BndrMap (mapTM (filterTM f) tm)
+ftBndrMap f (BndrMap tm) = BndrMap (fmap (filterTM f) tm)
 
 --------- Variable occurrence -------------
 data VarMap a = VM { vm_bvar   :: BoundVarMap a  -- Bound variable
                    , vm_fvar   :: DVarEnv a }      -- Free variable
 
+-- TODO(22292): derive
+instance Functor VarMap where
+    fmap f VM { vm_bvar = bv, vm_fvar = fv } = VM { vm_bvar = fmap f bv, vm_fvar = fmap f fv }
+
 instance TrieMap VarMap where
    type Key VarMap = Var
    emptyTM  = VM { vm_bvar = IntMap.empty, vm_fvar = emptyDVarEnv }
    lookupTM = lkVar emptyCME
    alterTM  = xtVar emptyCME
    foldTM   = fdVar
-   mapTM    = mapVar
    filterTM = ftVar
 
-mapVar :: (a->b) -> VarMap a -> VarMap b
-mapVar f (VM { vm_bvar = bv, vm_fvar = fv })
-  = VM { vm_bvar = mapTM f bv, vm_fvar = mapTM f fv }
-
 lkVar :: CmEnv -> Var -> VarMap a -> Maybe a
 lkVar env v
   | Just bv <- lookupCME env v = vm_bvar >.> lookupTM bv