module TyCoTidy
(
tidyType, tidyTypes,
tidyOpenType, tidyOpenTypes,
tidyOpenKind,
tidyVarBndr, tidyVarBndrs, tidyFreeTyCoVars, avoidNameClashes,
tidyOpenTyCoVar, tidyOpenTyCoVars,
tidyTyCoVarOcc,
tidyTopType,
tidyKind,
tidyCo, tidyCos,
tidyTyCoVarBinder, tidyTyCoVarBinders
) where
import GhcPrelude
import TyCoRep
import TyCoFVs (tyCoVarsOfTypesWellScoped, tyCoVarsOfTypeList)
import Name hiding (varName)
import Var
import VarEnv
import Util (seqList)
import Data.List (mapAccumL)
tidyVarBndrs :: TidyEnv -> [TyCoVar] -> (TidyEnv, [TyCoVar])
tidyVarBndrs tidy_env tvs
= mapAccumL tidyVarBndr (avoidNameClashes tvs tidy_env) tvs
tidyVarBndr :: TidyEnv -> TyCoVar -> (TidyEnv, TyCoVar)
tidyVarBndr tidy_env@(occ_env, subst) var
= case tidyOccName occ_env (getHelpfulOccName var) of
(occ_env', occ') -> ((occ_env', subst'), var')
where
subst' = extendVarEnv subst var var'
var' = setVarType (setVarName var name') type'
type' = tidyType tidy_env (varType var)
name' = tidyNameOcc name occ'
name = varName var
avoidNameClashes :: [TyCoVar] -> TidyEnv -> TidyEnv
avoidNameClashes tvs (occ_env, subst)
= (avoidClashesOccEnv occ_env occs, subst)
where
occs = map getHelpfulOccName tvs
getHelpfulOccName :: TyCoVar -> OccName
getHelpfulOccName tv
| isSystemName name, isTcTyVar tv
= mkTyVarOcc (occNameString occ ++ "0")
| otherwise
= occ
where
name = varName tv
occ = getOccName name
tidyTyCoVarBinder :: TidyEnv -> VarBndr TyCoVar vis
-> (TidyEnv, VarBndr TyCoVar vis)
tidyTyCoVarBinder tidy_env (Bndr tv vis)
= (tidy_env', Bndr tv' vis)
where
(tidy_env', tv') = tidyVarBndr tidy_env tv
tidyTyCoVarBinders :: TidyEnv -> [VarBndr TyCoVar vis]
-> (TidyEnv, [VarBndr TyCoVar vis])
tidyTyCoVarBinders tidy_env tvbs
= mapAccumL tidyTyCoVarBinder
(avoidNameClashes (binderVars tvbs) tidy_env) tvbs
tidyFreeTyCoVars :: TidyEnv -> [TyCoVar] -> TidyEnv
tidyFreeTyCoVars (full_occ_env, var_env) tyvars
= fst (tidyOpenTyCoVars (full_occ_env, var_env) tyvars)
tidyOpenTyCoVars :: TidyEnv -> [TyCoVar] -> (TidyEnv, [TyCoVar])
tidyOpenTyCoVars env tyvars = mapAccumL tidyOpenTyCoVar env tyvars
tidyOpenTyCoVar :: TidyEnv -> TyCoVar -> (TidyEnv, TyCoVar)
tidyOpenTyCoVar env@(_, subst) tyvar
= case lookupVarEnv subst tyvar of
Just tyvar' -> (env, tyvar')
Nothing ->
let env' = tidyFreeTyCoVars env (tyCoVarsOfTypeList (tyVarKind tyvar))
in tidyVarBndr env' tyvar
tidyTyCoVarOcc :: TidyEnv -> TyCoVar -> TyCoVar
tidyTyCoVarOcc env@(_, subst) tv
= case lookupVarEnv subst tv of
Nothing -> updateVarType (tidyType env) tv
Just tv' -> tv'
tidyTypes :: TidyEnv -> [Type] -> [Type]
tidyTypes env tys = map (tidyType env) tys
tidyType :: TidyEnv -> Type -> Type
tidyType _ (LitTy n) = LitTy n
tidyType env (TyVarTy tv) = TyVarTy (tidyTyCoVarOcc env tv)
tidyType env (TyConApp tycon tys) = let args = tidyTypes env tys
in args `seqList` TyConApp tycon args
tidyType env (AppTy fun arg) = (AppTy $! (tidyType env fun)) $! (tidyType env arg)
tidyType env ty@(FunTy _ arg res) = let { !arg' = tidyType env arg
; !res' = tidyType env res }
in ty { ft_arg = arg', ft_res = res' }
tidyType env (ty@(ForAllTy{})) = mkForAllTys' (zip tvs' vis) $! tidyType env' body_ty
where
(tvs, vis, body_ty) = splitForAllTys' ty
(env', tvs') = tidyVarBndrs env tvs
tidyType env (CastTy ty co) = (CastTy $! tidyType env ty) $! (tidyCo env co)
tidyType env (CoercionTy co) = CoercionTy $! (tidyCo env co)
mkForAllTys' :: [(TyCoVar, ArgFlag)] -> Type -> Type
mkForAllTys' tvvs ty = foldr strictMkForAllTy ty tvvs
where
strictMkForAllTy (tv,vis) ty = (ForAllTy $! ((Bndr $! tv) $! vis)) $! ty
splitForAllTys' :: Type -> ([TyCoVar], [ArgFlag], Type)
splitForAllTys' ty = go ty [] []
where
go (ForAllTy (Bndr tv vis) ty) tvs viss = go ty (tv:tvs) (vis:viss)
go ty tvs viss = (reverse tvs, reverse viss, ty)
tidyOpenTypes :: TidyEnv -> [Type] -> (TidyEnv, [Type])
tidyOpenTypes env tys
= (env', tidyTypes (trimmed_occ_env, var_env) tys)
where
(env'@(_, var_env), tvs') = tidyOpenTyCoVars env $
tyCoVarsOfTypesWellScoped tys
trimmed_occ_env = initTidyOccEnv (map getOccName tvs')
tidyOpenType :: TidyEnv -> Type -> (TidyEnv, Type)
tidyOpenType env ty = let (env', [ty']) = tidyOpenTypes env [ty] in
(env', ty')
tidyTopType :: Type -> Type
tidyTopType ty = tidyType emptyTidyEnv ty
tidyOpenKind :: TidyEnv -> Kind -> (TidyEnv, Kind)
tidyOpenKind = tidyOpenType
tidyKind :: TidyEnv -> Kind -> Kind
tidyKind = tidyType
tidyCo :: TidyEnv -> Coercion -> Coercion
tidyCo env@(_, subst) co
= go co
where
go_mco MRefl = MRefl
go_mco (MCo co) = MCo (go co)
go (Refl ty) = Refl (tidyType env ty)
go (GRefl r ty mco) = GRefl r (tidyType env ty) $! go_mco mco
go (TyConAppCo r tc cos) = let args = map go cos
in args `seqList` TyConAppCo r tc args
go (AppCo co1 co2) = (AppCo $! go co1) $! go co2
go (ForAllCo tv h co) = ((ForAllCo $! tvp) $! (go h)) $! (tidyCo envp co)
where (envp, tvp) = tidyVarBndr env tv
go (FunCo r co1 co2) = (FunCo r $! go co1) $! go co2
go (CoVarCo cv) = case lookupVarEnv subst cv of
Nothing -> CoVarCo cv
Just cv' -> CoVarCo cv'
go (HoleCo h) = HoleCo h
go (AxiomInstCo con ind cos) = let args = map go cos
in args `seqList` AxiomInstCo con ind args
go (UnivCo p r t1 t2) = (((UnivCo $! (go_prov p)) $! r) $!
tidyType env t1) $! tidyType env t2
go (SymCo co) = SymCo $! go co
go (TransCo co1 co2) = (TransCo $! go co1) $! go co2
go (NthCo r d co) = NthCo r d $! go co
go (LRCo lr co) = LRCo lr $! go co
go (InstCo co ty) = (InstCo $! go co) $! go ty
go (KindCo co) = KindCo $! go co
go (SubCo co) = SubCo $! go co
go (AxiomRuleCo ax cos) = let cos1 = tidyCos env cos
in cos1 `seqList` AxiomRuleCo ax cos1
go_prov UnsafeCoerceProv = UnsafeCoerceProv
go_prov (PhantomProv co) = PhantomProv (go co)
go_prov (ProofIrrelProv co) = ProofIrrelProv (go co)
go_prov p@(PluginProv _) = p
tidyCos :: TidyEnv -> [Coercion] -> [Coercion]
tidyCos env = map (tidyCo env)