{-# 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 "HsVersions.h"

-- | Update CafInfos and LFInfos of all occurrences (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 cg_infos mod_details =
  let
    ModDetails{ md_types = type_env -- for unfoldings
              , md_insts = insts
              , md_rules = rules
              } = mod_details

    -- type TypeEnv = NameEnv TyThing
    type_env' = mapNameEnv (updateTyThingIdInfos type_env' cg_infos) type_env
    -- NB: Knot-tied! The result, type_env', is passed right back into into
    -- updateTyThingIdInfos, so that that occurrences of any Ids (e.g. in
    -- IdInfos, etc) can be looked up in the tidied env

    !insts' = strictMap (updateInstIdInfos type_env' cg_infos) insts
    !rules' = strictMap (updateRuleIdInfos type_env') rules
  in
    mod_details{ md_types = type_env'
               , md_insts = insts'
               , md_rules = rules'
               }

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

updateRuleIdInfos :: TypeEnv -> CoreRule -> CoreRule
updateRuleIdInfos _ rule@BuiltinRule{} = rule
updateRuleIdInfos type_env Rule{ .. } = Rule { ru_rhs = updateGlobalIds type_env ru_rhs, .. }

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

updateInstIdInfos :: TypeEnv -> CgInfos -> ClsInst -> ClsInst
updateInstIdInfos type_env cg_infos =
    updateClsInstDFun (updateIdUnfolding type_env . updateIdInfo cg_infos)

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

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

updateTyThingIdInfos type_env cg_infos (AnId id) =
    AnId (updateIdUnfolding type_env (updateIdInfo cg_infos id))

updateTyThingIdInfos _ _ other = other -- AConLike, ATyCon, ACoAxiom

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

updateIdUnfolding :: TypeEnv -> Id -> Id
updateIdUnfolding type_env id =
    case idUnfolding id of
      CoreUnfolding{ .. } ->
        setIdUnfolding id CoreUnfolding{ uf_tmpl = updateGlobalIds type_env uf_tmpl, .. }
      DFunUnfolding{ .. } ->
        setIdUnfolding id DFunUnfolding{ df_args = map (updateGlobalIds type_env) df_args, .. }
      _ -> id

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

updateIdInfo :: CgInfos -> Id -> Id
updateIdInfo CgInfos{ cgNonCafs = NonCaffySet non_cafs, cgLFInfos = lf_infos } id =
    let
      not_caffy = elemNameSet (idName id) non_cafs
      mb_lf_info = lookupNameEnv lf_infos (idName id)

      id1 = if not_caffy then setIdCafInfo id NoCafRefs else id
      id2 = case mb_lf_info of
              Nothing -> id1
              Just lf_info -> setIdLFInfo id1 lf_info
    in
      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 env e = go env e
  where
    go_id :: NameEnv TyThing -> Id -> Id
    go_id env var =
      case lookupNameEnv env (varName var) of
        Nothing -> var
        Just (AnId id) -> id
        Just other -> pprPanic "UpdateIdInfos.updateGlobalIds" $
          text "Found a non-Id for Id Name" <+> ppr (varName var) $$
          nest 4 (text "Id:" <+> ppr var $$
                  text "TyThing:" <+> ppr other)

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

    go_binds :: NameEnv TyThing -> CoreBind -> CoreBind
    go_binds env (NonRec b e) =
      assertNotInNameEnv env [b] (NonRec b (go env e))
    go_binds env (Rec prs) =
      assertNotInNameEnv env (map fst prs) (Rec (mapSnd (go env) 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 env ids x = ASSERT(not (any (\id -> elemNameEnv (idName id) env) ids)) x