+type LevelEnv = (Bool, -- True <=> Float lambdas too
+ VarEnv Level, -- Domain is *post-cloned* TyVars and Ids
+ Subst, -- Domain is pre-cloned Ids; tracks the in-scope set
+ -- so that subtitution is capture-avoiding
+ IdEnv ([Var], LevelledExpr)) -- Domain is pre-cloned Ids
+ -- We clone let-bound variables so that they are still
+ -- distinct when floated out; hence the SubstEnv/IdEnv.
+ -- (see point 3 of the module overview comment).
+ -- We also use these envs when making a variable polymorphic
+ -- because we want to float it out past a big lambda.
+ --
+ -- The SubstEnv and IdEnv always implement the same mapping, but the
+ -- SubstEnv maps to CoreExpr and the IdEnv to LevelledExpr
+ -- Since the range is always a variable or type application,
+ -- there is never any difference between the two, but sadly
+ -- the types differ. The SubstEnv is used when substituting in
+ -- a variable's IdInfo; the IdEnv when we find a Var.
+ --
+ -- In addition the IdEnv records a list of tyvars free in the
+ -- type application, just so we don't have to call freeVars on
+ -- the type application repeatedly.
+ --
+ -- The domain of the both envs is *pre-cloned* Ids, though
+ --
+ -- The domain of the VarEnv Level is the *post-cloned* Ids
+
+initialEnv :: Bool -> LevelEnv
+initialEnv float_lams = (float_lams, emptyVarEnv, emptySubst, emptyVarEnv)
+
+floatLams :: LevelEnv -> Bool
+floatLams (float_lams, _, _, _) = float_lams
+
+extendLvlEnv :: LevelEnv -> [(Var,Level)] -> LevelEnv
+-- Used when *not* cloning
+extendLvlEnv (float_lams, lvl_env, subst, id_env) prs
+ = (float_lams,
+ foldl add_lvl lvl_env prs,
+ foldl del_subst subst prs,
+ foldl del_id id_env prs)
+ where
+ add_lvl env (v,l) = extendVarEnv env v l
+ del_subst env (v,_) = extendInScope env v
+ del_id env (v,_) = delVarEnv env v
+ -- We must remove any clone for this variable name in case of
+ -- shadowing. This bit me in the following case
+ -- (in nofib/real/gg/Spark.hs):
+ --
+ -- case ds of wild {
+ -- ... -> case e of wild {
+ -- ... -> ... wild ...
+ -- }
+ -- }
+ --
+ -- The inside occurrence of @wild@ was being replaced with @ds@,
+ -- incorrectly, because the SubstEnv was still lying around. Ouch!
+ -- KSW 2000-07.
+
+-- extendCaseBndrLvlEnv adds the mapping case-bndr->scrut-var if it can
+-- (see point 4 of the module overview comment)
+extendCaseBndrLvlEnv (float_lams, lvl_env, subst, id_env) (Var scrut_var) case_bndr lvl
+ = (float_lams,
+ extendVarEnv lvl_env case_bndr lvl,
+ extendSubst subst case_bndr (DoneEx (Var scrut_var)),
+ extendVarEnv id_env case_bndr ([scrut_var], Var scrut_var))
+
+extendCaseBndrLvlEnv env scrut case_bndr lvl
+ = extendLvlEnv env [(case_bndr,lvl)]
+
+extendPolyLvlEnv dest_lvl (float_lams, lvl_env, subst, id_env) abs_vars bndr_pairs
+ = (float_lams,
+ foldl add_lvl lvl_env bndr_pairs,
+ foldl add_subst subst bndr_pairs,
+ foldl add_id id_env bndr_pairs)
+ where
+ add_lvl env (v,v') = extendVarEnv env v' dest_lvl
+ add_subst env (v,v') = extendSubst env v (DoneEx (mkVarApps (Var v') abs_vars))
+ add_id env (v,v') = extendVarEnv env v ((v':abs_vars), mkVarApps (Var v') abs_vars)
+
+extendCloneLvlEnv lvl (float_lams, lvl_env, _, id_env) new_subst bndr_pairs
+ = (float_lams,
+ foldl add_lvl lvl_env bndr_pairs,
+ new_subst,
+ foldl add_id id_env bndr_pairs)
+ where
+ add_lvl env (v,v') = extendVarEnv env v' lvl
+ add_id env (v,v') = extendVarEnv env v ([v'], Var v')
+
+
+maxIdLevel :: LevelEnv -> VarSet -> Level
+maxIdLevel (_, lvl_env,_,id_env) var_set
+ = foldVarSet max_in tOP_LEVEL var_set
+ where
+ max_in in_var lvl = foldr max_out lvl (case lookupVarEnv id_env in_var of
+ Just (abs_vars, _) -> abs_vars
+ Nothing -> [in_var])
+
+ max_out out_var lvl
+ | isId out_var = case lookupVarEnv lvl_env out_var of
+ Just lvl' -> maxLvl lvl' lvl
+ Nothing -> lvl
+ | otherwise = lvl -- Ignore tyvars in *maxIdLevel*
+
+lookupVar :: LevelEnv -> Id -> LevelledExpr
+lookupVar (_, _, _, id_env) v = case lookupVarEnv id_env v of
+ Just (_, expr) -> expr
+ other -> Var v
+
+absVarsOf :: Level -> LevelEnv -> Var -> [Var]
+ -- If f is free in the exression, and f maps to poly_f a b c in the
+ -- current substitution, then we must report a b c as candidate type
+ -- variables
+absVarsOf dest_lvl (_, lvl_env, _, id_env) v
+ | isId v
+ = [final_av | av <- lookup_avs v, abstract_me av, final_av <- add_tyvars av]
+
+ | otherwise
+ = if abstract_me v then [v] else []
+
+ where
+ abstract_me v = case lookupVarEnv lvl_env v of
+ Just lvl -> dest_lvl `ltLvl` lvl
+ Nothing -> False
+
+ lookup_avs v = case lookupVarEnv id_env v of
+ Just (abs_vars, _) -> abs_vars
+ Nothing -> [v]
+
+ -- We are going to lambda-abstract, so nuke any IdInfo,
+ -- and add the tyvars of the Id
+ add_tyvars v | isId v = zap v : varSetElems (idFreeTyVars v)
+ | otherwise = [v]
+
+ zap v = WARN( workerExists (idWorkerInfo v)
+ || not (isEmptyCoreRules (idSpecialisation v)),
+ text "absVarsOf: discarding info on" <+> ppr v )
+ setIdInfo v vanillaIdInfo
+\end{code}
+
+\begin{code}