%
% (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, tcEqExpr, mkPiTypes )
import CoreFVs 		( exprsFreeVars )
import CoreTidy		( tidyRules )
import PprCore		( pprRules )
import WwLib		( mkWorkerArgs )
import DataCon		( dataConRepArity )
import Type		( tyConAppArgs )
import Id		( Id, idName, idType, isDataConWorkId_maybe, 
			  mkUserLocal, mkSysLocal )
import Var		( Var )
import VarEnv
import VarSet
import Name		( nameOccName, nameSrcLoc )
import Rules		( addIdSpecialisations, mkLocalRule, rulesOfBinds )
import OccName		( mkSpecOcc )
import ErrUtils		( dumpIfSet_dyn )
import DynFlags	( DynFlags, DynFlag(..) )
import BasicTypes	( Activation(..) )
import Outputable

import Maybes		( orElse )
import Util		( mapAccumL, lengthAtLeast, notNull )
import List		( nubBy, partition )
import UniqSupply
import Outputable
import FastString
\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"
		  (pprRules (tidyRules emptyTidyEnv (rulesOfBinds 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')
\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 ty 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 ty 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)])
  | notNull 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,
		       call_args `lengthAtLeast` 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 tcEqExpr 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, rule_number)
  = 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
	spec_body    = mkApps rhs pats
	body_ty	     = exprType spec_body
	
	(spec_lam_args, spec_call_args) = mkWorkerArgs bndrs body_ty
		-- Usual w/w hack to avoid generating 
		-- a spec_rhs of unlifted type and no args
	
	rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
	spec_rhs  = mkLams spec_lam_args spec_body
	spec_id   = mkUserLocal spec_occ spec_uniq (mkPiTypes spec_lam_args body_ty) fn_loc
	rule_rhs  = mkVarApps (Var spec_id) spec_call_args
	rule      = mkLocalRule rule_name specConstrActivation fn_name bndrs pats rule_rhs
    in
    returnUs (rule, (spec_id, spec_rhs))

-- In which phase should the specialise-constructor rules be active?
-- Originally I made them always-active, but Manuel found that
-- this defeated some clever user-written rules.  So Plan B
-- is to make them active only in Phase 0; after all, currently,
-- the specConstr transformation is only run after the simplifier
-- has reached Phase 0.  In general one would want it to be 
-- flag-controllable, but for now I'm leaving it baked in
--					[SLPJ Oct 01]
specConstrActivation :: Activation
specConstrActivation = ActiveAfter 0	-- Baked in; see comments above
\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 FSLIT("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 <- isDataConWorkId_maybe fun,
			  args `lengthAtLeast` 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}