path: root/ghc/compiler/specialise/SpecConstr.lhs

%
% (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
%
\section[SpecConstr]{Specialise over constructors}

\begin{code}
module SpecConstr(
	specConstrProgram	
    ) where

#include "HsVersions.h"

import CoreSyn
import CoreLint		( showPass, endPass )
import CoreUtils	( exprType, eqExpr )
import CoreFVs 		( exprsFreeVars )
import DataCon		( dataConRepArity )
import Type		( tyConAppArgs )
import PprCore		( pprCoreRules )
import Id		( Id, idName, idType, idSpecialisation,
			  isDataConId_maybe,
			  mkUserLocal, mkSysLocal )
import Var		( Var )
import VarEnv
import VarSet
import Name		( nameOccName, nameSrcLoc )
import Rules		( addIdSpecialisations )
import OccName		( mkSpecOcc )
import ErrUtils		( dumpIfSet_dyn )
import CmdLineOpts	( DynFlags, DynFlag(..) )
import BasicTypes	( Activation(..) )
import Outputable

import Maybes		( orElse )
import Util		( mapAccumL )
import List		( nubBy, partition )
import UniqSupply
import Outputable
\end{code}

-----------------------------------------------------
			Game plan
-----------------------------------------------------

Consider
	drop n []     = []
	drop 0 xs     = []
	drop n (x:xs) = drop (n-1) xs

After the first time round, we could pass n unboxed.  This happens in
numerical code too.  Here's what it looks like in Core:

	drop n xs = case xs of
		      []     -> []
		      (y:ys) -> case n of 
				  I# n# -> case n# of
					     0 -> []
					     _ -> drop (I# (n# -# 1#)) xs

Notice that the recursive call has an explicit constructor as argument.
Noticing this, we can make a specialised version of drop
	
	RULE: drop (I# n#) xs ==> drop' n# xs

	drop' n# xs = let n = I# n# in ...orig RHS...

Now the simplifier will apply the specialisation in the rhs of drop', giving

	drop' n# xs = case xs of
		      []     -> []
		      (y:ys) -> case n# of
				  0 -> []
				  _ -> drop (n# -# 1#) xs

Much better!  

We'd also like to catch cases where a parameter is carried along unchanged,
but evaluated each time round the loop:

	f i n = if i>0 || i>n then i else f (i*2) n

Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
In Core, by the time we've w/wd (f is strict in i) we get

	f i# n = case i# ># 0 of
		   False -> I# i#
		   True  -> case n of n' { I# n# ->
			    case i# ># n# of
				False -> I# i#
				True  -> f (i# *# 2#) n'

At the call to f, we see that the argument, n is know to be (I# n#),
and n is evaluated elsewhere in the body of f, so we can play the same
trick as above.  However we don't want to do that if the boxed version
of n is needed (else we'd avoid the eval but pay more for re-boxing n).
So in this case we want that the *only* uses of n are in case statements.


So we look for

* A self-recursive function.  Ignore mutual recursion for now, 
  because it's less common, and the code is simpler for self-recursion.

* EITHER

   a) At a recursive call, one or more parameters is an explicit 
      constructor application
	AND
      That same parameter is scrutinised by a case somewhere in 
      the RHS of the function

  OR

    b) At a recursive call, one or more parameters has an unfolding
       that is an explicit constructor application
	AND
      That same parameter is scrutinised by a case somewhere in 
      the RHS of the function
	AND
      Those are the only uses of the parameter


There's a bit of a complication with type arguments.  If the call
site looks like

	f p = ...f ((:) [a] x xs)...

then our specialised function look like

	f_spec x xs = let p = (:) [a] x xs in ....as before....

This only makes sense if either
  a) the type variable 'a' is in scope at the top of f, or
  b) the type variable 'a' is an argument to f (and hence fs)

Actually, (a) may hold for value arguments too, in which case
we may not want to pass them.  Supose 'x' is in scope at f's
defn, but xs is not.  Then we'd like

	f_spec xs = let p = (:) [a] x xs in ....as before....

Similarly (b) may hold too.  If x is already an argument at the
call, no need to pass it again.

Finally, if 'a' is not in scope at the call site, we could abstract
it as we do the term variables:

	f_spec a x xs = let p = (:) [a] x xs in ...as before...

So the grand plan is:

	* abstract the call site to a constructor-only pattern
	  e.g.  C x (D (f p) (g q))  ==>  C s1 (D s2 s3)

	* Find the free variables of the abstracted pattern

	* Pass these variables, less any that are in scope at
	  the fn defn.


NOTICE that we only abstract over variables that are not in scope,
so we're in no danger of shadowing variables used in "higher up"
in f_spec's RHS.


%************************************************************************
%*									*
\subsection{Top level wrapper stuff}
%*									*
%************************************************************************

\begin{code}
specConstrProgram :: DynFlags -> UniqSupply -> [CoreBind] -> IO [CoreBind]
specConstrProgram dflags us binds
  = do
	showPass dflags "SpecConstr"

	let (binds', _) = initUs us (go emptyScEnv binds)

	endPass dflags "SpecConstr" Opt_D_dump_spec binds'

	dumpIfSet_dyn dflags Opt_D_dump_rules "Top-level specialisations"
		  (vcat (map dump_specs (concat (map bindersOf binds'))))

	return binds'
  where
    go env []	        = returnUs []
    go env (bind:binds) = scBind env bind 	`thenUs` \ (env', _, bind') ->
			  go env' binds 	`thenUs` \ binds' ->
			  returnUs (bind' : binds')

dump_specs var = pprCoreRules var (idSpecialisation var)
\end{code}


%************************************************************************
%*									*
\subsection{Environment: goes downwards}
%*									*
%************************************************************************

\begin{code}
data ScEnv = SCE { scope :: VarEnv HowBound,
			-- Binds all non-top-level variables in scope

		   cons  :: ConstrEnv
	     }

type ConstrEnv = IdEnv (AltCon, [CoreArg])
	-- Variables known to be bound to a constructor
	-- in a particular case alternative

emptyScEnv = SCE { scope = emptyVarEnv, cons = emptyVarEnv }

data HowBound = RecFun		-- These are the recursive functions for which 
				-- we seek interesting call patterns

	      | RecArg		-- These are those functions' arguments; we are
				-- interested to see if those arguments are scrutinised

	      | Other		-- We track all others so we know what's in scope
				-- This is used in spec_one to check what needs to be
				-- passed as a parameter and what is in scope at the 
				-- function definition site

instance Outputable HowBound where
  ppr RecFun = text "RecFun"
  ppr RecArg = text "RecArg"
  ppr Other = text "Other"

lookupScopeEnv env v = lookupVarEnv (scope env) v

extendBndrs env bndrs = env { scope = extendVarEnvList (scope env) [(b,Other) | b <- bndrs] }
extendBndr  env bndr  = env { scope = extendVarEnv (scope env) bndr Other }

    -- When we encounter
    --	case scrut of b
    --	    C x y -> ...
    -- we want to bind b, and perhaps scrut too, to (C x y)
extendCaseBndrs :: ScEnv -> Id -> CoreExpr -> AltCon -> [Var] -> ScEnv
extendCaseBndrs env case_bndr scrut DEFAULT alt_bndrs
  = extendBndrs env (case_bndr : alt_bndrs)

extendCaseBndrs env case_bndr scrut con alt_bndrs
  = case scrut of
	Var v ->   -- Bind the scrutinee in the ConstrEnv if it's a variable
		   -- Also forget if the scrutinee is a RecArg, because we're
		   -- now in the branch of a case, and we don't want to
		   -- record a non-scrutinee use of v if we have
		   --	case v of { (a,b) -> ...(f v)... }
		 SCE { scope = extendVarEnv (scope env1) v Other,
		       cons  = extendVarEnv (cons env1)  v (con,args) }
	other -> env1

  where
    env1 = SCE { scope = extendVarEnvList (scope env) [(b,Other) | b <- case_bndr : alt_bndrs],
		 cons  = extendVarEnv     (cons  env) case_bndr (con,args) }

    args = map Type (tyConAppArgs (idType case_bndr)) ++
	   map varToCoreExpr alt_bndrs

    -- When we encounter a recursive function binding
    --	f = \x y -> ...
    -- we want to extend the scope env with bindings 
    -- that record that f is a RecFn and x,y are RecArgs
extendRecBndr env fn bndrs
  =  env { scope = scope env `extendVarEnvList` 
		   ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs]) }
\end{code}


%************************************************************************
%*									*
\subsection{Usage information: flows upwards}
%*									*
%************************************************************************

\begin{code}
data ScUsage
   = SCU {
	calls :: !(IdEnv ([Call])),	-- Calls
					-- The functions are a subset of the 
					-- 	RecFuns in the ScEnv

	occs :: !(IdEnv ArgOcc)		-- Information on argument occurrences
     }					-- The variables are a subset of the 
					--	RecArg in the ScEnv

type Call = (ConstrEnv, [CoreArg])
	-- The arguments of the call, together with the
	-- env giving the constructor bindings at the call site

nullUsage = SCU { calls = emptyVarEnv, occs = emptyVarEnv }

combineUsage u1 u2 = SCU { calls = plusVarEnv_C (++) (calls u1) (calls u2),
			   occs  = plusVarEnv_C combineOcc (occs u1) (occs u2) }

combineUsages [] = nullUsage
combineUsages us = foldr1 combineUsage us

data ArgOcc = CaseScrut 
	    | OtherOcc
	    | Both

instance Outputable ArgOcc where
  ppr CaseScrut = ptext SLIT("case-scrut")
  ppr OtherOcc  = ptext SLIT("other-occ")
  ppr Both      = ptext SLIT("case-scrut and other")

combineOcc CaseScrut CaseScrut = CaseScrut
combineOcc OtherOcc  OtherOcc  = OtherOcc
combineOcc _	     _	       = Both
\end{code}


%************************************************************************
%*									*
\subsection{The main recursive function}
%*									*
%************************************************************************

The main recursive function gathers up usage information, and
creates specialised versions of functions.

\begin{code}
scExpr :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
	-- The unique supply is needed when we invent
	-- a new name for the specialised function and its args

scExpr env e@(Type t) = returnUs (nullUsage, e)
scExpr env e@(Lit l)  = returnUs (nullUsage, e)
scExpr env e@(Var v)  = returnUs (varUsage env v OtherOcc, e)
scExpr env (Note n e) = scExpr env e	`thenUs` \ (usg,e') ->
			returnUs (usg, Note n e')
scExpr env (Lam b e)  = scExpr (extendBndr env b) e	`thenUs` \ (usg,e') ->
			returnUs (usg, Lam b e')

scExpr env (Case scrut b alts) 
  = sc_scrut scrut		`thenUs` \ (scrut_usg, scrut') ->
    mapAndUnzipUs sc_alt alts	`thenUs` \ (alts_usgs, alts') ->
    returnUs (combineUsages alts_usgs `combineUsage` scrut_usg,
	      Case scrut' b alts')
  where
    sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
    sc_scrut e	       = scExpr env e

    sc_alt (con,bs,rhs) = scExpr env1 rhs	`thenUs` \ (usg,rhs') ->
			  returnUs (usg, (con,bs,rhs'))
			where
			  env1 = extendCaseBndrs env b scrut con bs

scExpr env (Let bind body)
  = scBind env bind	`thenUs` \ (env', bind_usg, bind') ->
    scExpr env' body	`thenUs` \ (body_usg, body') ->
    returnUs (bind_usg `combineUsage` body_usg, Let bind' body')

scExpr env e@(App _ _) 
  = let 
	(fn, args) = collectArgs e
    in
    mapAndUnzipUs (scExpr env) args	`thenUs` \ (usgs, args') ->
    let
	arg_usg = combineUsages usgs
	fn_usg  | Var f <- fn,
		  Just RecFun <- lookupScopeEnv env f
		= SCU { calls = unitVarEnv f [(cons env, args)], 
			occs  = emptyVarEnv }
		| otherwise
		= nullUsage
    in
    returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
	-- Don't bother to look inside fn;
	-- it's almost always a variable

----------------------
scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
scBind env (Rec [(fn,rhs)])
  | not (null val_bndrs)
  = scExpr env_fn_body body		`thenUs` \ (usg, body') ->
    let
	SCU { calls = calls, occs = occs } = usg
    in
    specialise env fn bndrs body usg	`thenUs` \ (rules, spec_prs) ->
    returnUs (extendBndr env fn,	-- For the body of the letrec, just
					-- extend the env with Other to record 
					-- that it's in scope; no funny RecFun business
	      SCU { calls = calls `delVarEnv` fn, occs = occs `delVarEnvList` val_bndrs},
	      Rec ((fn `addIdSpecialisations` rules, mkLams bndrs body') : spec_prs))
  where
    (bndrs,body) = collectBinders rhs
    val_bndrs    = filter isId bndrs
    env_fn_body	 = extendRecBndr env fn bndrs

scBind env (Rec prs)
  = mapAndUnzipUs do_one prs	`thenUs` \ (usgs, prs') ->
    returnUs (extendBndrs env (map fst prs), combineUsages usgs, Rec prs')
  where
    do_one (bndr,rhs) = scExpr env rhs	`thenUs` \ (usg, rhs') ->
		        returnUs (usg, (bndr,rhs'))

scBind env (NonRec bndr rhs)
  = scExpr env rhs	`thenUs` \ (usg, rhs') ->
    returnUs (extendBndr env bndr, usg, NonRec bndr rhs')

----------------------
varUsage env v use 
  | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv, 
						occs = unitVarEnv v use }
  | otherwise		   	        = nullUsage
\end{code}


%************************************************************************
%*									*
\subsection{The specialiser}
%*									*
%************************************************************************

\begin{code}
specialise :: ScEnv
	   -> Id 			-- Functionn
	   -> [CoreBndr] -> CoreExpr	-- Its RHS
	   -> ScUsage			-- Info on usage
	   -> UniqSM ([CoreRule], 	-- Rules
		      [(Id,CoreExpr)])	-- Bindings

specialise env fn bndrs body (SCU {calls=calls, occs=occs})
  = getUs		`thenUs` \ us ->
    let
	all_calls = lookupVarEnv calls fn `orElse` []

	good_calls :: [[CoreArg]]
	good_calls = [ pats
		     | (con_env, call_args) <- all_calls,
		       length call_args >= n_bndrs,	    -- App is saturated
		       let call = (bndrs `zip` call_args),
		       any (good_arg con_env occs) call,    -- At least one arg is a constr app
		       let (_, pats) = argsToPats con_env us call_args
		     ]
    in
    mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
		  (nubBy same_call good_calls `zip` [1..])
  where
    n_bndrs  = length bndrs
    same_call as1 as2 = and (zipWith eqExpr as1 as2)

---------------------
good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
good_arg con_env arg_occs (bndr, arg)
  = case is_con_app_maybe con_env arg of	
	Just _ ->  bndr_usg_ok arg_occs bndr arg
	other   -> False

bndr_usg_ok :: IdEnv ArgOcc -> Var -> CoreArg -> Bool
bndr_usg_ok arg_occs bndr arg
  = case lookupVarEnv arg_occs bndr of
	Just CaseScrut -> True			-- Used only by case scrutiny
	Just Both      -> case arg of		-- Used by case and elsewhere
			    App _ _ -> True	-- so the arg should be an explicit con app
			    other   -> False
	other -> False				-- Not used, or used wonkily
    

---------------------
spec_one :: ScEnv
	 -> Id					-- Function
	 -> CoreExpr				-- Rhs of the original function
	 -> ([CoreArg], Int)
	 -> UniqSM (CoreRule, (Id,CoreExpr))	-- Rule and binding

-- spec_one creates a specialised copy of the function, together
-- with a rule for using it.  I'm very proud of how short this
-- function is, considering what it does :-).

{- 
  Example
  
     In-scope: a, x::a   
     f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
	  [c::*, v::(b,c) are presumably bound by the (...) part]
  ==>
     f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
		  (...entire RHS of f...) (b,c) ((:) (a,(b,c)) (x,v) hw)
  
     RULE:  forall b::* c::*,		-- Note, *not* forall a, x
		   v::(b,c),
		   hw::[(a,(b,c))] .
  
	    f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
-}

spec_one env fn rhs (pats, n)
  = getUniqueUs 			`thenUs` \ spec_uniq ->
    let 
	fn_name      = idName fn
	fn_loc       = nameSrcLoc fn_name
	spec_occ     = mkSpecOcc (nameOccName fn_name)
	pat_fvs	     = varSetElems (exprsFreeVars pats)
	vars_to_bind = filter not_avail pat_fvs
	not_avail v  = not (v `elemVarEnv` scope env)
		-- Put the type variables first; the type of a term
		-- variable may mention a type variable
	(tvs, ids)   = partition isTyVar vars_to_bind
	bndrs  	     = tvs ++ ids
	
	rule_name = _PK_ ("SC:" ++ showSDoc (ppr fn <> int n))
	spec_rhs  = mkLams bndrs (mkApps rhs pats)
	spec_id   = mkUserLocal spec_occ spec_uniq (exprType spec_rhs) fn_loc
	rule      = Rule rule_name AlwaysActive bndrs pats (mkVarApps (Var spec_id) bndrs)
    in
    returnUs (rule, (spec_id, spec_rhs))
\end{code}

%************************************************************************
%*									*
\subsection{Argument analysis}
%*									*
%************************************************************************

This code deals with analysing call-site arguments to see whether
they are constructor applications.

\begin{code}
    -- argToPat takes an actual argument, and returns an abstracted
    -- version, consisting of just the "constructor skeleton" of the
    -- argument, with non-constructor sub-expression replaced by new
    -- placeholder variables.  For example:
    --    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)

argToPat   :: ConstrEnv -> UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
argToPat env us (Type ty) 
  = (us, Type ty)

argToPat env us arg
  | Just (dc,args) <- is_con_app_maybe env arg
  = let
	(us',args') = argsToPats env us args
    in
    (us', mk_con_app dc args')

argToPat env us (Var v)	-- Don't uniqify existing vars,
  = (us, Var v)		-- so that we can spot when we pass them twice

argToPat env us arg
  = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
  where
    (us1,us2) = splitUniqSupply us

argsToPats :: ConstrEnv -> UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
argsToPats env us args = mapAccumL (argToPat env) us args
\end{code}


\begin{code}
is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe (AltCon, [CoreExpr])
is_con_app_maybe env (Var v)
  = lookupVarEnv env v
	-- You might think we could look in the idUnfolding here
	-- but that doesn't take account of which branch of a 
	-- case we are in, which is the whole point

is_con_app_maybe env (Lit lit)
  = Just (LitAlt lit, [])

is_con_app_maybe env expr
  = case collectArgs expr of
	(Var fun, args) | Just con <- isDataConId_maybe fun,
			  length args >= dataConRepArity con
		-- Might be > because the arity excludes type args
		        -> Just (DataAlt con,args)

	other -> Nothing

mk_con_app :: AltCon -> [CoreArg] -> CoreExpr
mk_con_app (LitAlt lit)  []   = Lit lit
mk_con_app (DataAlt con) args = mkConApp con args
\end{code}