{-# LANGUAGE CPP, BangPatterns, Strict, RecordWildCards #-}

module GHC.Iface.UpdateIdInfos
  ( updateModDetailsIdInfos
  ) where

import GHC.Prelude

import GHC.Core
import GHC.Core.InstEnv

import GHC.StgToCmm.Types (CgInfos (..))

import GHC.Types.Id
import GHC.Types.Id.Info
import GHC.Types.Name.Env
import GHC.Types.Name.Set
import GHC.Types.Var
import GHC.Types.TypeEnv
import GHC.Types.TyThing

import GHC.Unit.Module.ModDetails

import GHC.Utils.Misc
import GHC.Utils.Outputable
import GHC.Utils.Panic

#include "GhclibHsVersions.h"

-- | Update CafInfos and LFInfos of all occurences (in rules, unfoldings, class
-- instances).
--
-- See Note [Conveying CAF-info and LFInfo between modules] in
-- GHC.StgToCmm.Types.
updateModDetailsIdInfos
  :: CgInfos
  -> ModDetails -- ^ ModDetails to update
  -> ModDetails

updateModDetailsIdInfos :: CgInfos -> ModDetails -> ModDetails
updateModDetailsIdInfos CgInfos
cg_infos ModDetails
mod_details =
  let
    ModDetails{ md_types :: ModDetails -> TypeEnv
md_types = TypeEnv
type_env -- for unfoldings
              , md_insts :: ModDetails -> [ClsInst]
md_insts = [ClsInst]
insts
              , md_rules :: ModDetails -> [CoreRule]
md_rules = [CoreRule]
rules
              } = ModDetails
mod_details

    -- type TypeEnv = NameEnv TyThing
    ~TypeEnv
type_env' = (TyThing -> TyThing) -> TypeEnv -> TypeEnv
forall elt1 elt2. (elt1 -> elt2) -> NameEnv elt1 -> NameEnv elt2
mapNameEnv (TypeEnv -> CgInfos -> TyThing -> TyThing
updateTyThingIdInfos TypeEnv
type_env' CgInfos
cg_infos) TypeEnv
type_env
    -- Not strict!

    !insts' :: [ClsInst]
insts' = (ClsInst -> ClsInst) -> [ClsInst] -> [ClsInst]
forall a b. (a -> b) -> [a] -> [b]
strictMap (TypeEnv -> CgInfos -> ClsInst -> ClsInst
updateInstIdInfos TypeEnv
type_env' CgInfos
cg_infos) [ClsInst]
insts
    !rules' :: [CoreRule]
rules' = (CoreRule -> CoreRule) -> [CoreRule] -> [CoreRule]
forall a b. (a -> b) -> [a] -> [b]
strictMap (TypeEnv -> CoreRule -> CoreRule
updateRuleIdInfos TypeEnv
type_env') [CoreRule]
rules
  in
    ModDetails
mod_details{ md_types :: TypeEnv
md_types = TypeEnv
type_env'
               , md_insts :: [ClsInst]
md_insts = [ClsInst]
insts'
               , md_rules :: [CoreRule]
md_rules = [CoreRule]
rules'
               }

--------------------------------------------------------------------------------
-- Rules
--------------------------------------------------------------------------------

updateRuleIdInfos :: TypeEnv -> CoreRule -> CoreRule
updateRuleIdInfos :: TypeEnv -> CoreRule -> CoreRule
updateRuleIdInfos TypeEnv
_ rule :: CoreRule
rule@BuiltinRule{} = CoreRule
rule
updateRuleIdInfos TypeEnv
type_env Rule{ Bool
[Maybe Name]
[CoreBndr]
[CoreExpr]
Name
RuleName
Module
Activation
CoreExpr
IsOrphan
ru_rough :: CoreRule -> [Maybe Name]
ru_rhs :: CoreRule -> CoreExpr
ru_orphan :: CoreRule -> IsOrphan
ru_origin :: CoreRule -> Module
ru_name :: CoreRule -> RuleName
ru_local :: CoreRule -> Bool
ru_fn :: CoreRule -> Name
ru_bndrs :: CoreRule -> [CoreBndr]
ru_auto :: CoreRule -> Bool
ru_args :: CoreRule -> [CoreExpr]
ru_act :: CoreRule -> Activation
ru_local :: Bool
ru_orphan :: IsOrphan
ru_origin :: Module
ru_auto :: Bool
ru_rhs :: CoreExpr
ru_args :: [CoreExpr]
ru_bndrs :: [CoreBndr]
ru_rough :: [Maybe Name]
ru_fn :: Name
ru_act :: Activation
ru_name :: RuleName
.. } = Rule :: RuleName
-> Activation
-> Name
-> [Maybe Name]
-> [CoreBndr]
-> [CoreExpr]
-> CoreExpr
-> Bool
-> Module
-> IsOrphan
-> Bool
-> CoreRule
Rule { ru_rhs :: CoreExpr
ru_rhs = TypeEnv -> CoreExpr -> CoreExpr
updateGlobalIds TypeEnv
type_env CoreExpr
ru_rhs, Bool
[Maybe Name]
[CoreBndr]
[CoreExpr]
Name
RuleName
Module
Activation
IsOrphan
ru_rough :: [Maybe Name]
ru_orphan :: IsOrphan
ru_origin :: Module
ru_name :: RuleName
ru_local :: Bool
ru_fn :: Name
ru_bndrs :: [CoreBndr]
ru_auto :: Bool
ru_args :: [CoreExpr]
ru_act :: Activation
ru_local :: Bool
ru_orphan :: IsOrphan
ru_origin :: Module
ru_auto :: Bool
ru_args :: [CoreExpr]
ru_bndrs :: [CoreBndr]
ru_rough :: [Maybe Name]
ru_fn :: Name
ru_act :: Activation
ru_name :: RuleName
.. }

--------------------------------------------------------------------------------
-- Instances
--------------------------------------------------------------------------------

updateInstIdInfos :: TypeEnv -> CgInfos -> ClsInst -> ClsInst
updateInstIdInfos :: TypeEnv -> CgInfos -> ClsInst -> ClsInst
updateInstIdInfos TypeEnv
type_env CgInfos
cg_infos =
    (CoreBndr -> CoreBndr) -> ClsInst -> ClsInst
updateClsInstDFun (TypeEnv -> CoreBndr -> CoreBndr
updateIdUnfolding TypeEnv
type_env (CoreBndr -> CoreBndr)
-> (CoreBndr -> CoreBndr) -> CoreBndr -> CoreBndr
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CgInfos -> CoreBndr -> CoreBndr
updateIdInfo CgInfos
cg_infos)

--------------------------------------------------------------------------------
-- TyThings
--------------------------------------------------------------------------------

updateTyThingIdInfos :: TypeEnv -> CgInfos -> TyThing -> TyThing

updateTyThingIdInfos :: TypeEnv -> CgInfos -> TyThing -> TyThing
updateTyThingIdInfos TypeEnv
type_env CgInfos
cg_infos (AnId CoreBndr
id) =
    CoreBndr -> TyThing
AnId (TypeEnv -> CoreBndr -> CoreBndr
updateIdUnfolding TypeEnv
type_env (CgInfos -> CoreBndr -> CoreBndr
updateIdInfo CgInfos
cg_infos CoreBndr
id))

updateTyThingIdInfos TypeEnv
_ CgInfos
_ TyThing
other = TyThing
other -- AConLike, ATyCon, ACoAxiom

--------------------------------------------------------------------------------
-- Unfoldings
--------------------------------------------------------------------------------

updateIdUnfolding :: TypeEnv -> Id -> Id
updateIdUnfolding :: TypeEnv -> CoreBndr -> CoreBndr
updateIdUnfolding TypeEnv
type_env CoreBndr
id =
    case CoreBndr -> Unfolding
idUnfolding CoreBndr
id of
      CoreUnfolding{ Bool
CoreExpr
UnfoldingGuidance
UnfoldingSource
uf_tmpl :: Unfolding -> CoreExpr
uf_src :: Unfolding -> UnfoldingSource
uf_is_work_free :: Unfolding -> Bool
uf_is_value :: Unfolding -> Bool
uf_is_top :: Unfolding -> Bool
uf_is_conlike :: Unfolding -> Bool
uf_guidance :: Unfolding -> UnfoldingGuidance
uf_expandable :: Unfolding -> Bool
uf_guidance :: UnfoldingGuidance
uf_expandable :: Bool
uf_is_work_free :: Bool
uf_is_conlike :: Bool
uf_is_value :: Bool
uf_is_top :: Bool
uf_src :: UnfoldingSource
uf_tmpl :: CoreExpr
.. } ->
        CoreBndr -> Unfolding -> CoreBndr
setIdUnfolding CoreBndr
id CoreUnfolding :: CoreExpr
-> UnfoldingSource
-> Bool
-> Bool
-> Bool
-> Bool
-> Bool
-> UnfoldingGuidance
-> Unfolding
CoreUnfolding{ uf_tmpl :: CoreExpr
uf_tmpl = TypeEnv -> CoreExpr -> CoreExpr
updateGlobalIds TypeEnv
type_env CoreExpr
uf_tmpl, Bool
UnfoldingGuidance
UnfoldingSource
uf_src :: UnfoldingSource
uf_is_work_free :: Bool
uf_is_value :: Bool
uf_is_top :: Bool
uf_is_conlike :: Bool
uf_guidance :: UnfoldingGuidance
uf_expandable :: Bool
uf_guidance :: UnfoldingGuidance
uf_expandable :: Bool
uf_is_work_free :: Bool
uf_is_conlike :: Bool
uf_is_value :: Bool
uf_is_top :: Bool
uf_src :: UnfoldingSource
.. }
      DFunUnfolding{ [CoreBndr]
[CoreExpr]
DataCon
df_con :: Unfolding -> DataCon
df_bndrs :: Unfolding -> [CoreBndr]
df_args :: Unfolding -> [CoreExpr]
df_args :: [CoreExpr]
df_con :: DataCon
df_bndrs :: [CoreBndr]
.. } ->
        CoreBndr -> Unfolding -> CoreBndr
setIdUnfolding CoreBndr
id DFunUnfolding :: [CoreBndr] -> DataCon -> [CoreExpr] -> Unfolding
DFunUnfolding{ df_args :: [CoreExpr]
df_args = (CoreExpr -> CoreExpr) -> [CoreExpr] -> [CoreExpr]
forall a b. (a -> b) -> [a] -> [b]
map (TypeEnv -> CoreExpr -> CoreExpr
updateGlobalIds TypeEnv
type_env) [CoreExpr]
df_args, [CoreBndr]
DataCon
df_con :: DataCon
df_bndrs :: [CoreBndr]
df_con :: DataCon
df_bndrs :: [CoreBndr]
.. }
      Unfolding
_ -> CoreBndr
id

--------------------------------------------------------------------------------
-- Expressions
--------------------------------------------------------------------------------

updateIdInfo :: CgInfos -> Id -> Id
updateIdInfo :: CgInfos -> CoreBndr -> CoreBndr
updateIdInfo CgInfos{ cgNonCafs :: CgInfos -> NonCaffySet
cgNonCafs = NonCaffySet NameSet
non_cafs, cgLFInfos :: CgInfos -> ModuleLFInfos
cgLFInfos = ModuleLFInfos
lf_infos } CoreBndr
id =
    let
      not_caffy :: Bool
not_caffy = Name -> NameSet -> Bool
elemNameSet (CoreBndr -> Name
idName CoreBndr
id) NameSet
non_cafs
      mb_lf_info :: Maybe LambdaFormInfo
mb_lf_info = ModuleLFInfos -> Name -> Maybe LambdaFormInfo
forall a. NameEnv a -> Name -> Maybe a
lookupNameEnv ModuleLFInfos
lf_infos (CoreBndr -> Name
idName CoreBndr
id)

      id1 :: CoreBndr
id1 = if Bool
not_caffy then CoreBndr -> CafInfo -> CoreBndr
setIdCafInfo CoreBndr
id CafInfo
NoCafRefs else CoreBndr
id
      id2 :: CoreBndr
id2 = case Maybe LambdaFormInfo
mb_lf_info of
              Maybe LambdaFormInfo
Nothing -> CoreBndr
id1
              Just LambdaFormInfo
lf_info -> CoreBndr -> LambdaFormInfo -> CoreBndr
setIdLFInfo CoreBndr
id1 LambdaFormInfo
lf_info
    in
      CoreBndr
id2

--------------------------------------------------------------------------------

updateGlobalIds :: NameEnv TyThing -> CoreExpr -> CoreExpr
-- Update occurrences of GlobalIds as directed by 'env'
-- The 'env' maps a GlobalId to a version with accurate CAF info
-- (and in due course perhaps other back-end-related info)
updateGlobalIds :: TypeEnv -> CoreExpr -> CoreExpr
updateGlobalIds TypeEnv
env CoreExpr
e = TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e
  where
    go_id :: NameEnv TyThing -> Id -> Id
    go_id :: TypeEnv -> CoreBndr -> CoreBndr
go_id TypeEnv
env CoreBndr
var =
      case TypeEnv -> Name -> Maybe TyThing
forall a. NameEnv a -> Name -> Maybe a
lookupNameEnv TypeEnv
env (CoreBndr -> Name
varName CoreBndr
var) of
        Maybe TyThing
Nothing -> CoreBndr
var
        Just (AnId CoreBndr
id) -> CoreBndr
id
        Just TyThing
other -> String -> SDoc -> CoreBndr
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"UpdateIdInfos.updateGlobalIds" (SDoc -> CoreBndr) -> SDoc -> CoreBndr
forall a b. (a -> b) -> a -> b
$
          String -> SDoc
text String
"Found a non-Id for Id Name" SDoc -> SDoc -> SDoc
<+> Name -> SDoc
forall a. Outputable a => a -> SDoc
ppr (CoreBndr -> Name
varName CoreBndr
var) SDoc -> SDoc -> SDoc
$$
          Int -> SDoc -> SDoc
nest Int
4 (String -> SDoc
text String
"Id:" SDoc -> SDoc -> SDoc
<+> CoreBndr -> SDoc
forall a. Outputable a => a -> SDoc
ppr CoreBndr
var SDoc -> SDoc -> SDoc
$$
                  String -> SDoc
text String
"TyThing:" SDoc -> SDoc -> SDoc
<+> TyThing -> SDoc
forall a. Outputable a => a -> SDoc
ppr TyThing
other)

    go :: NameEnv TyThing -> CoreExpr -> CoreExpr
    go :: TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env (Var CoreBndr
v) = CoreBndr -> CoreExpr
forall b. CoreBndr -> Expr b
Var (TypeEnv -> CoreBndr -> CoreBndr
go_id TypeEnv
env CoreBndr
v)
    go TypeEnv
_ e :: CoreExpr
e@Lit{} = CoreExpr
e
    go TypeEnv
env (App CoreExpr
e1 CoreExpr
e2) = CoreExpr -> CoreExpr -> CoreExpr
forall b. Expr b -> Expr b -> Expr b
App (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e1) (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e2)
    go TypeEnv
env (Lam CoreBndr
b CoreExpr
e) = TypeEnv -> [CoreBndr] -> CoreExpr -> CoreExpr
forall a b. NameEnv a -> [CoreBndr] -> b -> b
assertNotInNameEnv TypeEnv
env [CoreBndr
b] (CoreBndr -> CoreExpr -> CoreExpr
forall b. b -> Expr b -> Expr b
Lam CoreBndr
b (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e))
    go TypeEnv
env (Let Bind CoreBndr
bs CoreExpr
e) = Bind CoreBndr -> CoreExpr -> CoreExpr
forall b. Bind b -> Expr b -> Expr b
Let (TypeEnv -> Bind CoreBndr -> Bind CoreBndr
go_binds TypeEnv
env Bind CoreBndr
bs) (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e)
    go TypeEnv
env (Case CoreExpr
e CoreBndr
b Type
ty [Alt CoreBndr]
alts) =
        TypeEnv -> [CoreBndr] -> CoreExpr -> CoreExpr
forall a b. NameEnv a -> [CoreBndr] -> b -> b
assertNotInNameEnv TypeEnv
env [CoreBndr
b] (CoreExpr -> CoreBndr -> Type -> [Alt CoreBndr] -> CoreExpr
forall b. Expr b -> b -> Type -> [Alt b] -> Expr b
Case (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e) CoreBndr
b Type
ty ((Alt CoreBndr -> Alt CoreBndr) -> [Alt CoreBndr] -> [Alt CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map Alt CoreBndr -> Alt CoreBndr
go_alt [Alt CoreBndr]
alts))
      where
         go_alt :: Alt CoreBndr -> Alt CoreBndr
go_alt (AltCon
k,[CoreBndr]
bs,CoreExpr
e) = TypeEnv -> [CoreBndr] -> Alt CoreBndr -> Alt CoreBndr
forall a b. NameEnv a -> [CoreBndr] -> b -> b
assertNotInNameEnv TypeEnv
env [CoreBndr]
bs (AltCon
k, [CoreBndr]
bs, TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e)
    go TypeEnv
env (Cast CoreExpr
e Coercion
c) = CoreExpr -> Coercion -> CoreExpr
forall b. Expr b -> Coercion -> Expr b
Cast (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e) Coercion
c
    go TypeEnv
env (Tick Tickish CoreBndr
t CoreExpr
e) = Tickish CoreBndr -> CoreExpr -> CoreExpr
forall b. Tickish CoreBndr -> Expr b -> Expr b
Tick Tickish CoreBndr
t (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e)
    go TypeEnv
_ e :: CoreExpr
e@Type{} = CoreExpr
e
    go TypeEnv
_ e :: CoreExpr
e@Coercion{} = CoreExpr
e

    go_binds :: NameEnv TyThing -> CoreBind -> CoreBind
    go_binds :: TypeEnv -> Bind CoreBndr -> Bind CoreBndr
go_binds TypeEnv
env (NonRec CoreBndr
b CoreExpr
e) =
      TypeEnv -> [CoreBndr] -> Bind CoreBndr -> Bind CoreBndr
forall a b. NameEnv a -> [CoreBndr] -> b -> b
assertNotInNameEnv TypeEnv
env [CoreBndr
b] (CoreBndr -> CoreExpr -> Bind CoreBndr
forall b. b -> Expr b -> Bind b
NonRec CoreBndr
b (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env CoreExpr
e))
    go_binds TypeEnv
env (Rec [(CoreBndr, CoreExpr)]
prs) =
      TypeEnv -> [CoreBndr] -> Bind CoreBndr -> Bind CoreBndr
forall a b. NameEnv a -> [CoreBndr] -> b -> b
assertNotInNameEnv TypeEnv
env (((CoreBndr, CoreExpr) -> CoreBndr)
-> [(CoreBndr, CoreExpr)] -> [CoreBndr]
forall a b. (a -> b) -> [a] -> [b]
map (CoreBndr, CoreExpr) -> CoreBndr
forall a b. (a, b) -> a
fst [(CoreBndr, CoreExpr)]
prs) ([(CoreBndr, CoreExpr)] -> Bind CoreBndr
forall b. [(b, Expr b)] -> Bind b
Rec ((CoreExpr -> CoreExpr)
-> [(CoreBndr, CoreExpr)] -> [(CoreBndr, CoreExpr)]
forall b c a. (b -> c) -> [(a, b)] -> [(a, c)]
mapSnd (TypeEnv -> CoreExpr -> CoreExpr
go TypeEnv
env) [(CoreBndr, CoreExpr)]
prs))

-- In `updateGlobaLIds` Names of local binders should not shadow Name of
-- globals. This assertion is to check that.
assertNotInNameEnv :: NameEnv a -> [Id] -> b -> b
assertNotInNameEnv :: NameEnv a -> [CoreBndr] -> b -> b
assertNotInNameEnv NameEnv a
env [CoreBndr]
ids b
x = ASSERT(not (any (\id -> elemNameEnv (idName id) env) ids)) x