{-# LANGUAGE CPP #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE FlexibleInstances #-} {-# LANGUAGE RankNTypes #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE UndecidableInstances #-} {- (c) The University of Glasgow 2006 (c) The GRASP/AQUA Project, Glasgow University, 1992-1998 -} {-# OPTIONS_GHC -Wno-orphans #-} -- Eq (DeBruijn CoreExpr) and Eq (DeBruijn CoreAlt) module GHC.Core.Map.Expr ( -- * Maps over Core expressions CoreMap, emptyCoreMap, extendCoreMap, lookupCoreMap, foldCoreMap, -- * 'TrieMap' class reexports TrieMap(..), insertTM, deleteTM, lkDFreeVar, xtDFreeVar, lkDNamed, xtDNamed, (>.>), (|>), (|>>), ) where #include "HsVersions.h" import GHC.Prelude import GHC.Data.TrieMap import GHC.Core.Map.Type import GHC.Core import GHC.Core.Type import GHC.Types.Tickish import GHC.Types.Var import GHC.Utils.Misc import GHC.Utils.Outputable import qualified Data.Map as Map import GHC.Types.Name.Env import Control.Monad( (>=>) ) {- This module implements TrieMaps over Core related data structures like CoreExpr or Type. It is built on the Tries from the TrieMap module. The code is very regular and boilerplate-like, but there is some neat handling of *binders*. In effect they are deBruijn numbered on the fly. -} ---------------------- -- Recall that -- Control.Monad.(>=>) :: (a -> Maybe b) -> (b -> Maybe c) -> a -> Maybe c -- The CoreMap makes heavy use of GenMap. However the CoreMap Types are not -- known when defining GenMap so we can only specialize them here. {-# SPECIALIZE lkG :: Key CoreMapX -> CoreMapG a -> Maybe a #-} {-# SPECIALIZE xtG :: Key CoreMapX -> XT a -> CoreMapG a -> CoreMapG a #-} {-# SPECIALIZE mapG :: (a -> b) -> CoreMapG a -> CoreMapG b #-} {-# SPECIALIZE fdG :: (a -> b -> b) -> CoreMapG a -> b -> b #-} {- ************************************************************************ * * CoreMap * * ************************************************************************ -} {- Note [Binders] ~~~~~~~~~~~~~~ * In general we check binders as late as possible because types are less likely to differ than expression structure. That's why cm_lam :: CoreMapG (TypeMapG a) rather than cm_lam :: TypeMapG (CoreMapG a) * We don't need to look at the type of some binders, notably - the case binder in (Case _ b _ _) - the binders in an alternative because they are totally fixed by the context Note [Empty case alternatives] ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ * For a key (Case e b ty (alt:alts)) we don't need to look the return type 'ty', because every alternative has that type. * For a key (Case e b ty []) we MUST look at the return type 'ty', because otherwise (Case (error () "urk") _ Int []) would compare equal to (Case (error () "urk") _ Bool []) which is utterly wrong (#6097) We could compare the return type regardless, but the wildly common case is that it's unnecessary, so we have two fields (cm_case and cm_ecase) for the two possibilities. Only cm_ecase looks at the type. See also Note [Empty case alternatives] in GHC.Core. -} -- | @CoreMap a@ is a map from 'CoreExpr' to @a@. If you are a client, this -- is the type you want. newtype CoreMap a = CoreMap (CoreMapG a) 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 -- key makes it suitable for recursive traversal, since it can track binders, -- but it is strictly internal to this module. If you are including a 'CoreMap' -- inside another 'TrieMap', this is the type you want. type CoreMapG = GenMap CoreMapX -- | @CoreMapX a@ is the base map from @DeBruijn CoreExpr@ to @a@, but without -- the 'GenMap' optimization. data CoreMapX a = CM { cm_var :: VarMap a , cm_lit :: LiteralMap a , cm_co :: CoercionMapG a , cm_type :: TypeMapG a , cm_cast :: CoreMapG (CoercionMapG a) , cm_tick :: CoreMapG (TickishMap a) , cm_app :: CoreMapG (CoreMapG a) , cm_lam :: CoreMapG (BndrMap a) -- Note [Binders] , cm_letn :: CoreMapG (CoreMapG (BndrMap a)) , cm_letr :: ListMap CoreMapG (CoreMapG (ListMap BndrMap a)) , cm_case :: CoreMapG (ListMap AltMap a) , cm_ecase :: CoreMapG (TypeMapG a) -- Note [Empty case alternatives] } instance Eq (DeBruijn CoreExpr) where D env1 e1 == D env2 e2 = go e1 e2 where go (Var v1) (Var v2) = case (lookupCME env1 v1, lookupCME env2 v2) of (Just b1, Just b2) -> b1 == b2 (Nothing, Nothing) -> v1 == v2 _ -> False go (Lit lit1) (Lit lit2) = lit1 == lit2 go (Type t1) (Type t2) = D env1 t1 == D env2 t2 go (Coercion co1) (Coercion co2) = D env1 co1 == D env2 co2 go (Cast e1 co1) (Cast e2 co2) = D env1 co1 == D env2 co2 && go e1 e2 go (App f1 a1) (App f2 a2) = go f1 f2 && go a1 a2 -- This seems a bit dodgy, see 'eqTickish' go (Tick n1 e1) (Tick n2 e2) = n1 == n2 && go e1 e2 go (Lam b1 e1) (Lam b2 e2) = D env1 (varType b1) == D env2 (varType b2) && D env1 (varMultMaybe b1) == D env2 (varMultMaybe b2) && D (extendCME env1 b1) e1 == D (extendCME env2 b2) e2 go (Let (NonRec v1 r1) e1) (Let (NonRec v2 r2) e2) = go r1 r2 && D (extendCME env1 v1) e1 == D (extendCME env2 v2) e2 go (Let (Rec ps1) e1) (Let (Rec ps2) e2) = equalLength ps1 ps2 && D env1' rs1 == D env2' rs2 && D env1' e1 == D env2' e2 where (bs1,rs1) = unzip ps1 (bs2,rs2) = unzip ps2 env1' = extendCMEs env1 bs1 env2' = extendCMEs env2 bs2 go (Case e1 b1 t1 a1) (Case e2 b2 t2 a2) | null a1 -- See Note [Empty case alternatives] = null a2 && go e1 e2 && D env1 t1 == D env2 t2 | otherwise = go e1 e2 && D (extendCME env1 b1) a1 == D (extendCME env2 b2) a2 go _ _ = False emptyE :: CoreMapX a emptyE = CM { cm_var = emptyTM, cm_lit = emptyTM , cm_co = emptyTM, cm_type = emptyTM , cm_cast = emptyTM, cm_app = emptyTM , cm_lam = emptyTM, cm_letn = emptyTM , cm_letr = emptyTM, cm_case = emptyTM , cm_ecase = emptyTM, cm_tick = emptyTM } 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 , 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 = 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 } -------------------------- lookupCoreMap :: CoreMap a -> CoreExpr -> Maybe a lookupCoreMap cm e = lookupTM e cm extendCoreMap :: CoreMap a -> CoreExpr -> a -> CoreMap a extendCoreMap m e v = alterTM e (\_ -> Just v) m foldCoreMap :: (a -> b -> b) -> b -> CoreMap a -> b foldCoreMap k z m = foldTM k m z emptyCoreMap :: CoreMap a emptyCoreMap = emptyTM instance Outputable a => Outputable (CoreMap a) where ppr m = text "CoreMap elts" <+> ppr (foldTM (:) m []) ------------------------- fdE :: (a -> b -> b) -> CoreMapX a -> b -> b fdE k m = foldTM k (cm_var m) . foldTM k (cm_lit m) . foldTM k (cm_co m) . foldTM k (cm_type m) . foldTM (foldTM k) (cm_cast m) . foldTM (foldTM k) (cm_tick m) . foldTM (foldTM k) (cm_app m) . foldTM (foldTM k) (cm_lam m) . foldTM (foldTM (foldTM k)) (cm_letn m) . foldTM (foldTM (foldTM k)) (cm_letr m) . foldTM (foldTM k) (cm_case m) . foldTM (foldTM k) (cm_ecase m) -- lkE: lookup in trie for expressions lkE :: DeBruijn CoreExpr -> CoreMapX a -> Maybe a lkE (D env expr) cm = go expr cm where go (Var v) = cm_var >.> lkVar env v go (Lit l) = cm_lit >.> lookupTM l go (Type t) = cm_type >.> lkG (D env t) go (Coercion c) = cm_co >.> lkG (D env c) go (Cast e c) = cm_cast >.> lkG (D env e) >=> lkG (D env c) go (Tick tickish e) = cm_tick >.> lkG (D env e) >=> lkTickish tickish go (App e1 e2) = cm_app >.> lkG (D env e2) >=> lkG (D env e1) go (Lam v e) = cm_lam >.> lkG (D (extendCME env v) e) >=> lkBndr env v go (Let (NonRec b r) e) = cm_letn >.> lkG (D env r) >=> lkG (D (extendCME env b) e) >=> lkBndr env b go (Let (Rec prs) e) = let (bndrs,rhss) = unzip prs env1 = extendCMEs env bndrs in cm_letr >.> lkList (lkG . D env1) rhss >=> lkG (D env1 e) >=> lkList (lkBndr env1) bndrs go (Case e b ty as) -- See Note [Empty case alternatives] | null as = cm_ecase >.> lkG (D env e) >=> lkG (D env ty) | otherwise = cm_case >.> lkG (D env e) >=> lkList (lkA (extendCME env b)) as xtE :: DeBruijn CoreExpr -> XT a -> CoreMapX a -> CoreMapX a xtE (D env (Var v)) f m = m { cm_var = cm_var m |> xtVar env v f } xtE (D env (Type t)) f m = m { cm_type = cm_type m |> xtG (D env t) f } xtE (D env (Coercion c)) f m = m { cm_co = cm_co m |> xtG (D env c) f } xtE (D _ (Lit l)) f m = m { cm_lit = cm_lit m |> alterTM l f } xtE (D env (Cast e c)) f m = m { cm_cast = cm_cast m |> xtG (D env e) |>> xtG (D env c) f } xtE (D env (Tick t e)) f m = m { cm_tick = cm_tick m |> xtG (D env e) |>> xtTickish t f } xtE (D env (App e1 e2)) f m = m { cm_app = cm_app m |> xtG (D env e2) |>> xtG (D env e1) f } xtE (D env (Lam v e)) f m = m { cm_lam = cm_lam m |> xtG (D (extendCME env v) e) |>> xtBndr env v f } xtE (D env (Let (NonRec b r) e)) f m = m { cm_letn = cm_letn m |> xtG (D (extendCME env b) e) |>> xtG (D env r) |>> xtBndr env b f } xtE (D env (Let (Rec prs) e)) f m = m { cm_letr = let (bndrs,rhss) = unzip prs env1 = extendCMEs env bndrs in cm_letr m |> xtList (xtG . D env1) rhss |>> xtG (D env1 e) |>> xtList (xtBndr env1) bndrs f } xtE (D env (Case e b ty as)) f m | null as = m { cm_ecase = cm_ecase m |> xtG (D env e) |>> xtG (D env ty) f } | otherwise = m { cm_case = cm_case m |> xtG (D env e) |>> let env1 = extendCME env b in xtList (xtA env1) as f } -- TODO: this seems a bit dodgy, see 'eqTickish' type TickishMap a = Map.Map CoreTickish a lkTickish :: CoreTickish -> TickishMap a -> Maybe a lkTickish = lookupTM xtTickish :: CoreTickish -> XT a -> TickishMap a -> TickishMap a xtTickish = alterTM ------------------------ data AltMap a -- A single alternative = AM { am_deflt :: CoreMapG a , am_data :: DNameEnv (CoreMapG a) , am_lit :: LiteralMap (CoreMapG a) } instance TrieMap AltMap where type Key AltMap = CoreAlt emptyTM = AM { am_deflt = emptyTM , am_data = emptyDNameEnv , am_lit = emptyTM } lookupTM = lkA emptyCME alterTM = xtA emptyCME foldTM = fdA mapTM = mapA filterTM = ftA instance Eq (DeBruijn CoreAlt) where D env1 a1 == D env2 a2 = go a1 a2 where go (Alt DEFAULT _ rhs1) (Alt DEFAULT _ rhs2) = D env1 rhs1 == D env2 rhs2 go (Alt (LitAlt lit1) _ rhs1) (Alt (LitAlt lit2) _ rhs2) = lit1 == lit2 && D env1 rhs1 == D env2 rhs2 go (Alt (DataAlt dc1) bs1 rhs1) (Alt (DataAlt dc2) bs2 rhs2) = dc1 == dc2 && 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 } lkA :: CmEnv -> CoreAlt -> AltMap a -> Maybe a lkA env (Alt DEFAULT _ rhs) = am_deflt >.> lkG (D env rhs) lkA env (Alt (LitAlt lit) _ rhs) = am_lit >.> lookupTM lit >=> lkG (D env rhs) lkA env (Alt (DataAlt dc) bs rhs) = am_data >.> lkDNamed dc >=> lkG (D (extendCMEs env bs) rhs) xtA :: CmEnv -> CoreAlt -> XT a -> AltMap a -> AltMap a xtA env (Alt DEFAULT _ rhs) f m = m { am_deflt = am_deflt m |> xtG (D env rhs) f } xtA env (Alt (LitAlt l) _ rhs) f m = m { am_lit = am_lit m |> alterTM l |>> xtG (D env rhs) f } xtA env (Alt (DataAlt d) bs rhs) f m = m { am_data = am_data m |> xtDNamed d |>> xtG (D (extendCMEs env bs) rhs) f } fdA :: (a -> b -> b) -> AltMap a -> b -> b fdA k m = foldTM k (am_deflt m) . foldTM (foldTM k) (am_data m) . foldTM (foldTM k) (am_lit m)