GHC/Core/Map/Expr.hs

    {-# 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)