[project @ 2001-10-30 10:57:50 by simonpj]
[ghc-hetmet.git] / ghc / compiler / specialise / SpecConstr.lhs
index 029ec17..574e039 100644 (file)
@@ -12,11 +12,14 @@ module SpecConstr(
 
 import CoreSyn
 import CoreLint                ( showPass, endPass )
-import CoreUtils       ( exprType, exprIsConApp_maybe, eqExpr )
+import CoreUtils       ( exprType, eqExpr )
 import CoreFVs                 ( exprsFreeVars )
-import DataCon         ( isExistentialDataCon )
+import DataCon         ( dataConRepArity )
+import Type            ( tyConAppArgs )
 import PprCore         ( pprCoreRules )
-import Id              ( Id, idName, idSpecialisation, mkUserLocal, mkSysLocal )
+import Id              ( Id, idName, idType, idSpecialisation,
+                         isDataConId_maybe,
+                         mkUserLocal, mkSysLocal )
 import Var             ( Var )
 import VarEnv
 import VarSet
@@ -25,10 +28,11 @@ 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 Util            ( mapAccumL, lengthAtLeast )
 import List            ( nubBy, partition )
 import UniqSupply
 import Outputable
@@ -191,14 +195,22 @@ dump_specs var = pprCoreRules var (idSpecialisation var)
 
 %************************************************************************
 %*                                                                     *
-\subsection{Environments and such}
+\subsection{Environment: goes downwards}
 %*                                                                     *
 %************************************************************************
 
 \begin{code}
-type ScEnv = VarEnv HowBound
+data ScEnv = SCE { scope :: VarEnv HowBound,
+                       -- Binds all non-top-level variables in scope
 
-emptyScEnv = emptyVarEnv
+                  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
@@ -211,19 +223,73 @@ data HowBound = RecFun            -- These are the recursive functions for which
                                -- passed as a parameter and what is in scope at the 
                                -- function definition site
 
-extendBndrs env bndrs = extendVarEnvList env [(b,Other) | b <- bndrs]
-extendBndr  env bndr  = extendVarEnv env bndr Other
+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 ([[CoreArg]])),        -- Calls
-                                               -- The functions are a subset of the 
-                                               --      RecFuns in the ScEnv
+       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),
@@ -253,6 +319,9 @@ combineOcc _             _         = Both
 %*                                                                     *
 %************************************************************************
 
+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
@@ -275,10 +344,10 @@ scExpr env (Case scrut b alts)
     sc_scrut e@(Var v) = returnUs (varUsage env v CaseScrut, e)
     sc_scrut e        = scExpr env e
 
-    sc_alt (con,bs,rhs) = scExpr env rhs       `thenUs` \ (usg,rhs') ->
+    sc_alt (con,bs,rhs) = scExpr env1 rhs      `thenUs` \ (usg,rhs') ->
                          returnUs (usg, (con,bs,rhs'))
                        where
-                         env1 = extendBndrs env (b:bs)
+                         env1 = extendCaseBndrs env b scrut con bs
 
 scExpr env (Let bind body)
   = scBind env bind    `thenUs` \ (env', bind_usg, bind') ->
@@ -293,8 +362,9 @@ scExpr env e@(App _ _)
     let
        arg_usg = combineUsages usgs
        fn_usg  | Var f <- fn,
-                 Just RecFun <- lookupVarEnv env f
-               = SCU { calls = unitVarEnv f [args], occs = emptyVarEnv }
+                 Just RecFun <- lookupScopeEnv env f
+               = SCU { calls = unitVarEnv f [(cons env, args)], 
+                       occs  = emptyVarEnv }
                | otherwise
                = nullUsage
     in
@@ -306,15 +376,20 @@ scExpr env e@(App _ _)
 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 scBind env (Rec [(fn,rhs)])
   | not (null val_bndrs)
-  = scExpr env' body                   `thenUs` \ (usg@(SCU { calls = calls, occs = occs }), body') ->
+  = 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 (extendBndrs env bndrs,
+    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' = env `extendVarEnvList` ((fn,RecFun): [(bndr,RecArg) | bndr <- bndrs])
+    env_fn_body         = extendRecBndr env fn bndrs
 
 scBind env (Rec prs)
   = mapAndUnzipUs do_one prs   `thenUs` \ (usgs, prs') ->
@@ -329,8 +404,9 @@ scBind env (NonRec bndr rhs)
 
 ----------------------
 varUsage env v use 
-  | Just RecArg <- lookupVarEnv env v = SCU { calls = emptyVarEnv, occs = unitVarEnv v use }
-  | otherwise                        = nullUsage
+  | Just RecArg <- lookupScopeEnv env v = SCU { calls = emptyVarEnv, 
+                                               occs = unitVarEnv v use }
+  | otherwise                          = nullUsage
 \end{code}
 
 
@@ -355,14 +431,13 @@ specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 
        good_calls :: [[CoreArg]]
        good_calls = [ pats
-                    | call_args <- all_calls,
-                      length call_args >= n_bndrs,     -- App is saturated
+                    | (con_env, call_args) <- all_calls,
+                      call_args `lengthAtLeast` n_bndrs,           -- App is saturated
                       let call = (bndrs `zip` call_args),
-                      any (good_arg occs) call,
-                      let (_, pats) = argsToPats us 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
-    pprTrace "specialise" (ppr all_calls $$ ppr good_calls) $
     mapAndUnzipUs (spec_one env fn (mkLams bndrs body)) 
                  (nubBy same_call good_calls `zip` [1..])
   where
@@ -370,17 +445,15 @@ specialise env fn bndrs body (SCU {calls=calls, occs=occs})
     same_call as1 as2 = and (zipWith eqExpr as1 as2)
 
 ---------------------
-good_arg :: IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
-good_arg arg_occs (bndr, arg)
-  = case exprIsConApp_maybe arg of                     -- exprIsConApp_maybe looks
-       Just (dc,_) -> not (isExistentialDataCon dc)    -- through unfoldings
-                      && bndr_usg_ok arg_occs bndr arg
+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
-  = pprTrace "bndr_ok" (ppr bndr <+> ppr (lookupVarEnv arg_occs bndr)) $
-    case lookupVarEnv arg_occs bndr of
+  = 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
@@ -389,52 +462,31 @@ bndr_usg_ok arg_occs bndr arg
     
 
 ---------------------
-argsToPats :: UniqSupply -> [CoreArg] -> (UniqSupply, [CoreExpr])
-argsToPats us args = mapAccumL argToPat us args
-
-argToPat   :: UniqSupply -> CoreArg   -> (UniqSupply, CoreExpr)
---    C a (D (f x) (g y))  ==>  C p1 (D p2 p3)
-argToPat us (Type ty) 
-  = (us, Type ty)
-
-argToPat us arg
-  | Just (dc,args) <- exprIsConApp_maybe arg
-  = let
-       (us',args') = argsToPats us args
-    in
-    (us', mkConApp dc args')
-
-argToPat us (Var v)    -- Don't uniqify existing vars,
-  = (us, Var v)                -- so that we can spot when we pass them twice
-
-argToPat us arg
-  = (us1, Var (mkSysLocal SLIT("sc") (uniqFromSupply us2) (exprType arg)))
-  where
-    (us1,us2) = splitUniqSupply us
-
----------------------
 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)) v (h v))...
-         [c is presumably bound by the (...) part]
+     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::(a,(b,c)) -> 
-                 (...entire RHS of f...) (b,c) ((:) (a,(b,c)) v (h v))
+     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,
-                  y::[(a,(b,c))], 
-                  v::(a,(b,c)), 
-                  h::(a,(b,c))->[(a,(b,c))] .
+     RULE:  forall b::* c::*,          -- Note, *not* forall a, x
+                  v::(b,c),
+                  hw::[(a,(b,c))] .
   
-           f (b,c) ((:) (a,(b,c)) v (h v)) = f_spec b c v
+           f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
 -}
 
 spec_one env fn rhs (pats, n)
@@ -445,15 +497,93 @@ spec_one env fn rhs (pats, n)
        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` env)
-               -- Put the type variables first just for tidiness
+       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 pat_fvs pats (mkVarApps (Var spec_id) bndrs)
+       rule      = Rule rule_name specConstrActivation
+                        bndrs pats (mkVarApps (Var spec_id) bndrs)
     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 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,
+                         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}