Commit 7aa3f524 authored by simonpj@microsoft.com's avatar simonpj@microsoft.com
Browse files

Make the LiberateCase transformation understand associated types

Consider this FC program:
	data family AT a :: *
	data instance AT Int = T1 Int Int

	f :: AT Int -> Int
	f t = case t of DEFAULT -> <body>

We'd like to replace the DEFAULT by a use of T1, so that if 
we scrutinise t inside <body> we share the evaluation:

	f t = case (t `cast` co) of T1 x y -> <body>

I decided to do this as part of the liberate-case transformation,
which is already trying to avoid redundant evals.  

The new transformation requires knowledge of the family instance
environment, so I had to extend ModGuts to carry the fam_inst_env,
and put that envt into the liberate-case environment. 

Otherwise it's all pretty straightforward.
parent 899fd7fb
......@@ -60,23 +60,24 @@ deSugar :: HscEnv -> ModLocation -> TcGblEnv -> IO (Maybe ModGuts)
deSugar hsc_env
mod_loc
tcg_env@(TcGblEnv { tcg_mod = mod,
tcg_src = hsc_src,
tcg_type_env = type_env,
tcg_imports = imports,
tcg_exports = exports,
tcg_dus = dus,
tcg_inst_uses = dfun_uses_var,
tcg_th_used = th_var,
tcg_keep = keep_var,
tcg_rdr_env = rdr_env,
tcg_fix_env = fix_env,
tcg_deprecs = deprecs,
tcg_binds = binds,
tcg_fords = fords,
tcg_rules = rules,
tcg_insts = insts,
tcg_fam_insts = fam_insts })
tcg_env@(TcGblEnv { tcg_mod = mod,
tcg_src = hsc_src,
tcg_type_env = type_env,
tcg_imports = imports,
tcg_exports = exports,
tcg_dus = dus,
tcg_inst_uses = dfun_uses_var,
tcg_th_used = th_var,
tcg_keep = keep_var,
tcg_rdr_env = rdr_env,
tcg_fix_env = fix_env,
tcg_fam_inst_env = fam_inst_env,
tcg_deprecs = deprecs,
tcg_binds = binds,
tcg_fords = fords,
tcg_rules = rules,
tcg_insts = insts,
tcg_fam_insts = fam_insts })
= do { showPass dflags "Desugar"
-- Desugar the program
......@@ -156,23 +157,24 @@ deSugar hsc_env
-- sort to get into canonical order
mod_guts = ModGuts {
mg_module = mod,
mg_boot = isHsBoot hsc_src,
mg_exports = exports,
mg_deps = deps,
mg_usages = usages,
mg_dir_imps = [m | (m,_,_) <- moduleEnvElts dir_imp_mods],
mg_rdr_env = rdr_env,
mg_fix_env = fix_env,
mg_deprecs = deprecs,
mg_types = type_env,
mg_insts = insts,
mg_fam_insts = fam_insts,
mg_rules = ds_rules,
mg_binds = ds_binds,
mg_foreign = ds_fords,
mg_hpc_info = ds_hpc_info,
mg_dbg_sites = dbgSites }
mg_module = mod,
mg_boot = isHsBoot hsc_src,
mg_exports = exports,
mg_deps = deps,
mg_usages = usages,
mg_dir_imps = [m | (m,_,_) <- moduleEnvElts dir_imp_mods],
mg_rdr_env = rdr_env,
mg_fix_env = fix_env,
mg_deprecs = deprecs,
mg_types = type_env,
mg_insts = insts,
mg_fam_insts = fam_insts,
mg_fam_inst_env = fam_inst_env,
mg_rules = ds_rules,
mg_binds = ds_binds,
mg_foreign = ds_fords,
mg_hpc_info = ds_hpc_info,
mg_dbg_sites = dbgSites }
; return (Just mod_guts)
}}}
......
......@@ -485,7 +485,10 @@ data ModGuts
mg_rdr_env :: !GlobalRdrEnv, -- Top-level lexical environment
mg_fix_env :: !FixityEnv, -- Fixity env, for things declared in
-- this module
mg_deprecs :: !Deprecations, -- Deprecations declared in the module
mg_fam_inst_env :: FamInstEnv, -- Type-family instance enviroment
-- for *home-package* modules (including
-- this one). c.f. tcg_fam_inst_env
mg_types :: !TypeEnv,
mg_insts :: ![Instance], -- Instances
......@@ -493,6 +496,7 @@ data ModGuts
mg_rules :: ![CoreRule], -- Rules from this module
mg_binds :: ![CoreBind], -- Bindings for this module
mg_foreign :: !ForeignStubs,
mg_deprecs :: !Deprecations, -- Deprecations declared in the module
mg_hpc_info :: !HpcInfo, -- info about coverage tick boxes
mg_dbg_sites :: ![(SiteNumber, Coord)] -- Bkpts inserted by the renamer
}
......
......@@ -8,18 +8,28 @@ module LiberateCase ( liberateCase ) where
#include "HsVersions.h"
import DynFlags ( DynFlags, DynFlag(..) )
import StaticFlags ( opt_LiberateCaseThreshold )
import DynFlags
import HscTypes
import CoreLint ( showPass, endPass )
import CoreSyn
import CoreUnfold ( couldBeSmallEnoughToInline )
import Id ( Id, setIdName, idName, setIdNotExported )
import Rules ( RuleBase )
import UniqSupply ( UniqSupply )
import SimplMonad ( SimplCount, zeroSimplCount )
import Id
import FamInstEnv
import Type
import Coercion
import TyCon
import VarEnv
import Name ( localiseName )
import Outputable
import Util ( notNull )
import Data.IORef ( readIORef )
\end{code}
The liberate-case transformation
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
This module walks over @Core@, and looks for @case@ on free variables.
The criterion is:
if there is case on a free on the route to the recursive call,
......@@ -27,30 +37,24 @@ The criterion is:
Example
\begin{verbatim}
f = \ t -> case v of
V a b -> a : f t
\end{verbatim}
f = \ t -> case v of
V a b -> a : f t
=> the inner f is replaced.
\begin{verbatim}
f = \ t -> case v of
V a b -> a : (letrec
f = \ t -> case v of
V a b -> a : (letrec
f = \ t -> case v of
V a b -> a : f t
in f) t
\end{verbatim}
in f) t
(note the NEED for shadowing)
=> Simplify
\begin{verbatim}
f = \ t -> case v of
V a b -> a : (letrec
f = \ t -> case v of
V a b -> a : (letrec
f = \ t -> a : f t
in f t)
\begin{verbatim}
in f t)
Better code, because 'a' is free inside the inner letrec, rather
than needing projection from v.
......@@ -72,7 +76,6 @@ We'd like to avoid the redundant pattern match, transforming to
(is this necessarily an improvement)
Similarly drop:
drop n [] = []
......@@ -119,66 +122,64 @@ scope. For example:
Here, the level of @f@ is zero, the level of @g@ is one,
and the level of @h@ is zero (NB not one).
\begin{code}
type LibCaseLevel = Int
topLevel :: LibCaseLevel
topLevel = 0
\end{code}
\begin{code}
data LibCaseEnv
= LibCaseEnv {
lc_size :: Int, -- Bomb-out size for deciding if
-- potential liberatees are too big.
-- (passed in from cmd-line args)
lc_lvl :: LibCaseLevel, -- Current level
lc_lvl_env :: IdEnv LibCaseLevel,
-- Binds all non-top-level in-scope Ids
-- (top-level and imported things have
-- a level of zero)
lc_rec_env :: IdEnv CoreBind,
-- Binds *only* recursively defined ids,
-- to their own binding group,
-- and *only* in their own RHSs
Note [Indexed data types]
~~~~~~~~~~~~~~~~~~~~~~~~~
Consider
data family T :: * -> *
data T Int = TI Int
f :: T Int -> Bool
f x = case x of { DEFAULT -> <body> }
We would like to change this to
f x = case x `cast` co of { TI p -> <body> }
so that <body> can make use of the fact that x is already evaluated to
a TI; and a case on a known data type may be more efficient than a
polymorphic one (not sure this is true any longer). Anyway the former
showed up in Roman's experiments. Example:
foo :: FooT Int -> Int -> Int
foo t n = t `seq` bar n
where
bar 0 = 0
bar n = bar (n - case t of TI i -> i)
Here we'd like to avoid repeated evaluating t inside the loop, by
taking advantage of the `seq`.
We implement this as part of the liberate-case transformation by
spotting
case <scrut> of (x::T) tys { DEFAULT -> <body> }
where x :: T tys, and T is a indexed family tycon. Find the
representation type (T77 tys'), and coercion co, and transform to
case <scrut> `cast` co of (y::T77 tys')
DEFAULT -> let x = y `cast` sym co in <body>
The "find the representation type" part is done by looking up in the
family-instance environment.
NB: in fact we re-use x (changing its type) to avoid making a fresh y;
this entails shadowing, but that's ok.
%************************************************************************
%* *
Top-level code
%* *
%************************************************************************
lc_scruts :: [(Id,LibCaseLevel)]
-- Each of these Ids was scrutinised by an
-- enclosing case expression, with the
-- specified number of enclosing
-- recursive bindings; furthermore,
-- the Id is bound at a lower level
-- than the case expression. The order is
-- insignificant; it's a bag really
-- lc_fams :: FamInstEnvs
-- Instance env for indexed data types
}
initEnv :: Int -> LibCaseEnv
initEnv bomb_size
= LibCaseEnv { lc_size = bomb_size, lc_lvl = 0,
lc_lvl_env = emptyVarEnv, lc_rec_env = emptyVarEnv,
lc_scruts = [] }
bombOutSize = lc_size
\end{code}
Programs
~~~~~~~~
\begin{code}
liberateCase :: DynFlags -> [CoreBind] -> IO [CoreBind]
liberateCase dflags binds
= do {
showPass dflags "Liberate case" ;
let { binds' = do_prog (initEnv opt_LiberateCaseThreshold) binds } ;
endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
{- no specific flag for dumping -}
}
liberateCase :: HscEnv -> UniqSupply -> RuleBase -> ModGuts
-> IO (SimplCount, ModGuts)
liberateCase hsc_env _ _ guts
= do { let dflags = hsc_dflags hsc_env
; eps <- readIORef (hsc_EPS hsc_env)
; let fam_envs = (eps_fam_inst_env eps, mg_fam_inst_env guts)
; showPass dflags "Liberate case"
; let { env = initEnv dflags fam_envs
; binds' = do_prog env (mg_binds guts) }
; endPass dflags "Liberate case" Opt_D_verbose_core2core binds'
{- no specific flag for dumping -}
; return (zeroSimplCount dflags, guts { mg_binds = binds' }) }
where
do_prog env [] = []
do_prog env (bind:binds) = bind' : do_prog env' binds
......@@ -186,9 +187,15 @@ liberateCase dflags binds
(env', bind') = libCaseBind env bind
\end{code}
%************************************************************************
%* *
Main payload
%* *
%************************************************************************
Bindings
~~~~~~~~
\begin{code}
libCaseBind :: LibCaseEnv -> CoreBind -> (LibCaseEnv, CoreBind)
......@@ -254,7 +261,7 @@ libCase env (Let bind body)
(env_body, bind') = libCaseBind env bind
libCase env (Case scrut bndr ty alts)
= Case (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
= mkCase env (libCase env scrut) bndr ty (map (libCaseAlt env_alts) alts)
where
env_alts = addBinders (mk_alt_env scrut) [bndr]
mk_alt_env (Var scrut_var) = addScrutedVar env scrut_var
......@@ -264,6 +271,24 @@ libCase env (Case scrut bndr ty alts)
libCaseAlt env (con,args,rhs) = (con, args, libCase (addBinders env args) rhs)
\end{code}
\begin{code}
mkCase :: LibCaseEnv -> CoreExpr -> Id -> Type -> [CoreAlt] -> CoreExpr
-- See Note [Indexed data types]
mkCase env scrut bndr ty [(DEFAULT,_,rhs)]
| Just (tycon, tys) <- splitTyConApp_maybe (idType bndr)
, [(subst, fam_inst)] <- lookupFamInstEnv (lc_fams env) tycon tys
= let
rep_tc = famInstTyCon fam_inst
rep_tys = map (substTyVar subst) (tyConTyVars rep_tc)
bndr' = setIdType bndr (mkTyConApp rep_tc rep_tys)
Just co_tc = tyConFamilyCoercion_maybe rep_tc
co = mkTyConApp co_tc rep_tys
bind = NonRec bndr (Cast (Var bndr') (mkSymCoercion co))
in mkCase env (Cast scrut co) bndr' ty [(DEFAULT,[],Let bind rhs)]
mkCase env scrut bndr ty alts
= Case scrut bndr ty alts
\end{code}
Ids
~~~
\begin{code}
......@@ -282,9 +307,12 @@ libCaseId env v
\end{code}
%************************************************************************
%* *
Utility functions
%* *
%************************************************************************
Utility functions
~~~~~~~~~~~~~~~~~
\begin{code}
addBinders :: LibCaseEnv -> [CoreBndr] -> LibCaseEnv
addBinders env@(LibCaseEnv { lc_lvl = lvl, lc_lvl_env = lvl_env }) binders
......@@ -335,3 +363,62 @@ freeScruts :: LibCaseEnv
freeScruts env rec_bind_lvl
= [v | (v,scrut_lvl) <- lc_scruts env, scrut_lvl > rec_bind_lvl]
\end{code}
%************************************************************************
%* *
The environment
%* *
%************************************************************************
\begin{code}
type LibCaseLevel = Int
topLevel :: LibCaseLevel
topLevel = 0
\end{code}
\begin{code}
data LibCaseEnv
= LibCaseEnv {
lc_size :: Int, -- Bomb-out size for deciding if
-- potential liberatees are too big.
-- (passed in from cmd-line args)
lc_lvl :: LibCaseLevel, -- Current level
lc_lvl_env :: IdEnv LibCaseLevel,
-- Binds all non-top-level in-scope Ids
-- (top-level and imported things have
-- a level of zero)
lc_rec_env :: IdEnv CoreBind,
-- Binds *only* recursively defined ids,
-- to their own binding group,
-- and *only* in their own RHSs
lc_scruts :: [(Id,LibCaseLevel)],
-- Each of these Ids was scrutinised by an
-- enclosing case expression, with the
-- specified number of enclosing
-- recursive bindings; furthermore,
-- the Id is bound at a lower level
-- than the case expression. The order is
-- insignificant; it's a bag really
lc_fams :: FamInstEnvs
-- Instance env for indexed data types
}
initEnv :: DynFlags -> FamInstEnvs -> LibCaseEnv
initEnv dflags fams
= LibCaseEnv { lc_size = libCaseThreshold dflags,
lc_lvl = 0,
lc_lvl_env = emptyVarEnv,
lc_rec_env = emptyVarEnv,
lc_scruts = [],
lc_fams = fams }
bombOutSize = lc_size
\end{code}
......@@ -134,7 +134,7 @@ doCorePasses hsc_env rb us stats guts (to_do : to_dos)
doCorePass (CoreDoSimplify mode sws) = _scc_ "Simplify" simplifyPgm mode sws
doCorePass CoreCSE = _scc_ "CommonSubExpr" trBinds cseProgram
doCorePass CoreLiberateCase = _scc_ "LiberateCase" trBinds liberateCase
doCorePass CoreLiberateCase = _scc_ "LiberateCase" liberateCase
doCorePass CoreDoFloatInwards = _scc_ "FloatInwards" trBinds floatInwards
doCorePass (CoreDoFloatOutwards f) = _scc_ "FloatOutwards" trBindsU (floatOutwards f)
doCorePass CoreDoStaticArgs = _scc_ "StaticArgs" trBinds doStaticArgs
......@@ -148,6 +148,8 @@ doCorePass (CoreDoRuleCheck phase pat) = observe (ruleCheck phase pat)
doCorePass CoreDoNothing = observe (\ _ _ -> return ())
#ifdef OLD_STRICTNESS
doCorePass CoreDoOldStrictness = _scc_ "OldStrictness" trBinds doOldStrictness
#else
doCorePass CoreDoOldStrictness = panic "CoreDoOldStrictness"
#endif
#ifdef OLD_STRICTNESS
......
......@@ -302,6 +302,7 @@ tcRnExtCore hsc_env (HsExtCore this_mod decls src_binds)
mg_types = final_type_env,
mg_insts = tcg_insts tcg_env,
mg_fam_insts = tcg_fam_insts tcg_env,
mg_fam_inst_env = tcg_fam_inst_env tcg_env,
mg_rules = [],
mg_binds = core_binds,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment