content: {-# 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,
     -- * Alpha equality
     eqDeBruijnExpr, eqCoreExpr,
     -- * 'TrieMap' class reexports
     TrieMap(..), insertTM, deleteTM,
     lkDFreeVar, xtDFreeVar,
     lkDNamed, xtDNamed,
     (>.>), (|>), (|>>),
   ) where
  
  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( (>=>) )
  import GHC.Types.Literal (Literal)
  
  {-
  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)
  
  -- 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
      filterTM f (CoreMap m) = CoreMap (filterTM f m)
      mapMaybeTM f (CoreMap m) = CoreMap (mapMaybeTM 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
  
  type LiteralMap  a = Map.Map Literal a
  
  -- | @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
      (==) = eqDeBruijnExpr
  
  eqDeBruijnExpr :: DeBruijn CoreExpr -> DeBruijn CoreExpr -> Bool
  eqDeBruijnExpr (D env1 e1) (D env2 e2) = go e1 e2 where
      go (Var v1) (Var v2)             = eqDeBruijnVar (D env1 v1) (D env2 v2)
      go (Lit lit1)    (Lit lit2)      = lit1 == lit2
      go (Type t1)    (Type t2)        = eqDeBruijnType (D env1 t1) (D env2 t2)
      -- See Note [Alpha-equality for Coercion arguments]
      go (Coercion {}) (Coercion {}) = True
      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
      go (Tick n1 e1) (Tick n2 e2)
        =  eqDeBruijnTickish (D env1 n1) (D env2 n2)
        && go e1 e2
  
      go (Lam b1 e1)  (Lam b2 e2)
        =  eqDeBruijnType (D env1 (varType b1)) (D env2 (varType b2))
        && D env1 (varMultMaybe b1) == D env2 (varMultMaybe b2)
        && eqDeBruijnExpr (D (extendCME env1 b1) e1) (D (extendCME env2 b2) e2)
  
      go (Let (NonRec v1 r1) e1) (Let (NonRec v2 r2) e2)
        =  go r1 r2 -- See Note [Alpha-equality for let-bindings]
        && eqDeBruijnExpr (D (extendCME env1 v1) e1) (D (extendCME env2 v2) e2)
  
      go (Let (Rec ps1) e1) (Let (Rec ps2) e2)
        = equalLength ps1 ps2
        -- See Note [Alpha-equality for let-bindings]
        && all2 (\b1 b2 -> eqDeBruijnType (D env1 (varType b1))
                                          (D env2 (varType b2)))
                bs1 bs2
        && D env1' rs1 == D env2' rs2
        && eqDeBruijnExpr (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
  
  eqDeBruijnTickish :: DeBruijn CoreTickish -> DeBruijn CoreTickish -> Bool
  eqDeBruijnTickish (D env1 t1) (D env2 t2) = go t1 t2 where
      go (Breakpoint lext lid lids) (Breakpoint rext rid rids)
          =  lid == rid
          && D env1 lids == D env2 rids
          && lext == rext
      go l r = l == r
  
  -- Compares for equality, modulo alpha
  eqCoreExpr :: CoreExpr -> CoreExpr -> Bool
  eqCoreExpr e1 e2 = eqDeBruijnExpr (deBruijnize e1) (deBruijnize e2)
  
  {- Note [Alpha-equality for Coercion arguments]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  The 'Coercion' constructor only appears in argument positions, and so, if the
  functions are equal, then the arguments must have equal types. Because the
  comparison for coercions (correctly) checks only their types, checking for
  alpha-equality of the coercions is redundant.
  -}
  
  {- Note [Alpha-equality for let-bindings]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  For /recursive/ let-bindings we need to check that the types of the binders
  are alpha-equivalent. Otherwise
  
    letrec (x : Bool) = x in x
  
  and
  
    letrec (y : Char) = y in y
  
  would be considered alpha-equivalent, which they are obviously not.
  
  For /non-recursive/ let-bindings, we do not have to check that the types of
  the binders are alpha-equivalent. When the RHSs (the expressions) of the
  non-recursive let-binders are well-formed and well-typed (which we assume they
  are at this point in the compiler), and the RHSs are alpha-equivalent, then the
  bindings must have the same type.
  
  In addition, it is also worth pointing out that
  
    letrec { x = e1; y = e2 } in b
  
  is NOT considered equal to
  
    letrec { y = e2; x = e1 } in b
  -}
  
  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 }
  
  -- 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
     filterTM = ftE
     mapMaybeTM = mpE
  
  --------------------------
  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 = 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 }
  
  mpE :: (a -> Maybe b) -> CoreMapX a -> CoreMapX b
  mpE 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 = mapMaybeTM f cvar, cm_lit = mapMaybeTM f clit
         , cm_co = mapMaybeTM f cco, cm_type = mapMaybeTM f ctype
         , cm_cast = fmap (mapMaybeTM f) ccast, cm_app = fmap (mapMaybeTM f) capp
         , cm_lam = fmap (mapMaybeTM f) clam, cm_letn = fmap (fmap (mapMaybeTM f)) cletn
         , cm_letr = fmap (fmap (mapMaybeTM f)) cletr, cm_case = fmap (mapMaybeTM f) ccase
         , cm_ecase = fmap (mapMaybeTM f) cecase, cm_tick = fmap (mapMaybeTM 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) }
  
  -- 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
                   , am_data = emptyDNameEnv
                   , am_lit  = emptyTM }
     lookupTM = lkA emptyCME
     alterTM  = xtA emptyCME
     foldTM   = fdA
     filterTM = ftA
     mapMaybeTM = mpA
  
  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
  
  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 = 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)
  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)
  
  mpA :: (a -> Maybe b) -> AltMap a -> AltMap b
  mpA f (AM { am_deflt = adeflt, am_data = adata, am_lit = alit })
    = AM { am_deflt = mapMaybeTM f adeflt
         , am_data = fmap (mapMaybeTM f) adata
         , am_lit = fmap (mapMaybeTM f) alit }
  