Make SpecConstr work better for nested functions
[ghc-hetmet.git] / compiler / specialise / SpecConstr.lhs
index 74944da..9d1ba01 100644 (file)
@@ -22,7 +22,7 @@ import DataCon                ( dataConRepArity, isVanillaDataCon )
 import Type            ( tyConAppArgs, tyVarsOfTypes )
 import Unify           ( coreRefineTys )
 import Id              ( Id, idName, idType, isDataConWorkId_maybe, 
-                         mkUserLocal, mkSysLocal )
+                         mkUserLocal, mkSysLocal, idUnfolding )
 import Var             ( Var )
 import VarEnv
 import VarSet
@@ -98,6 +98,8 @@ 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.
 
 
+Note [Good arguments]
+~~~~~~~~~~~~~~~~~~~~~
 So we look for
 
 * A self-recursive function.  Ignore mutual recursion for now, 
@@ -122,6 +124,8 @@ So we look for
       Those are the only uses of the parameter
 
 
+What to abstract over
+~~~~~~~~~~~~~~~~~~~~~
 There's a bit of a complication with type arguments.  If the call
 site looks like
 
@@ -157,7 +161,7 @@ So the grand plan is:
        * Find the free variables of the abstracted pattern
 
        * Pass these variables, less any that are in scope at
-         the fn defn.
+         the fn defn.  But see Note [Shadowing] below.
 
 
 NOTICE that we only abstract over variables that are not in scope,
@@ -165,6 +169,30 @@ so we're in no danger of shadowing variables used in "higher up"
 in f_spec's RHS.
 
 
+Note [Shadowing]
+~~~~~~~~~~~~~~~~
+In this pass we gather up usage information that may mention variables
+that are bound between the usage site and the definition site; or (more
+seriously) may be bound to something different at the definition site.
+For example:
+
+       f x = letrec g y v = let x = ... 
+                            in ...(g (a,b) x)...
+
+Since 'x' is in scope at the call site, we may make a rewrite rule that 
+looks like
+       RULE forall a,b. g (a,b) x = ...
+But this rule will never match, because it's really a different 'x' at 
+the call site -- and that difference will be manifest by the time the
+simplifier gets to it.  [A worry: the simplifier doesn't *guarantee*
+no-shadowing, so perhaps it may not be distinct?]
+
+Anyway, the rule isn't actually wrong, it's just not useful.  One possibility
+is to run deShadowBinds before running SpecConstr, but instead we run the
+simplifier.  That gives the simplest possible program for SpecConstr to
+chew on; and it virtually guarantees no shadowing.
+
+
 %************************************************************************
 %*                                                                     *
 \subsection{Top level wrapper stuff}
@@ -211,6 +239,10 @@ data ConValue  = CV AltCon [CoreArg]
        -- Variables known to be bound to a constructor
        -- in a particular case alternative
 
+
+instance Outputable ConValue where
+   ppr (CV con args) = ppr con <+> interpp'SP args
+
 refineConstrEnv :: Subst -> ConstrEnv -> ConstrEnv
 -- The substitution is a type substitution only
 refineConstrEnv subst env = mapVarEnv refine_con_value env
@@ -262,6 +294,9 @@ extendCaseBndrs env case_bndr scrut con@(DataAlt data_con) alt_bndrs
                   map varToCoreExpr alt_bndrs
 
     gadt_args = map (substExpr subst . varToCoreExpr) alt_bndrs
+       -- This call generates some bogus warnings from substExpr,
+       -- because it's inconvenient to put all the Ids in scope
+       -- Will be fixed when we move to FC
 
     (alt_tvs, _) = span isTyVar alt_bndrs
     Just (tv_subst, is_local) = coreRefineTys data_con alt_tvs (idType case_bndr)
@@ -388,29 +423,33 @@ scExpr env e@(App _ _)
   = let 
        (fn, args) = collectArgs e
     in
-    mapAndUnzipUs (scExpr env) args    `thenUs` \ (usgs, args') ->
+    mapAndUnzipUs (scExpr env) (fn:args)       `thenUs` \ (usgs, (fn':args')) ->
+       -- Process the function too.   It's almost always a variable,
+       -- but not always.  In particular, if this pass follows float-in,
+       -- which it may, we can get 
+       --      (let f = ...f... in f) arg1 arg2
     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
+       call_usg = case fn of
+                       Var f | Just RecFun <- lookupScopeEnv env f
+                             -> SCU { calls = unitVarEnv f [(cons env, args)], 
+                                      occs  = emptyVarEnv }
+                       other -> nullUsage
     in
-    returnUs (arg_usg `combineUsage` fn_usg, mkApps fn args')
-       -- Don't bother to look inside fn;
-       -- it's almost always a variable
+    returnUs (combineUsages usgs `combineUsage` call_usg, mkApps fn' args')
+
 
 ----------------------
 scBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, ScUsage, CoreBind)
 scBind env (Rec [(fn,rhs)])
   | notNull val_bndrs
   = scExpr env_fn_body body            `thenUs` \ (usg, body') ->
+    specialise env fn bndrs body' usg  `thenUs` \ (rules, spec_prs) ->
+       -- Note body': the specialised copies should be based on the 
+       --             optimised version of the body, in case there were
+       --             nested functions inside.
     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
@@ -463,7 +502,7 @@ specialise env fn bndrs body (SCU {calls=calls, occs=occs})
        good_calls = [ pats
                     | (con_env, call_args) <- all_calls,
                       call_args `lengthAtLeast` n_bndrs,           -- App is saturated
-                      let call = (bndrs `zip` call_args),
+                      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
                     ]
@@ -476,6 +515,7 @@ specialise env fn bndrs body (SCU {calls=calls, occs=occs})
 
 ---------------------
 good_arg :: ConstrEnv -> IdEnv ArgOcc -> (CoreBndr, CoreArg) -> Bool
+-- See Note [Good arguments] above
 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
@@ -527,6 +567,8 @@ spec_one env fn rhs (pats, rule_number)
        spec_occ     = mkSpecOcc (nameOccName fn_name)
        pat_fvs      = varSetElems (exprsFreeVars pats)
        vars_to_bind = filter not_avail pat_fvs
+               -- See Note [Shadowing] at the top
+
        not_avail v  = not (v `elemVarEnv` scope env)
                -- Put the type variables first; the type of a term
                -- variable may mention a type variable
@@ -602,10 +644,20 @@ argsToPats env us args = mapAccumL (argToPat env) us args
 \begin{code}
 is_con_app_maybe :: ConstrEnv -> CoreExpr -> Maybe ConValue
 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
+  = case lookupVarEnv env v of
+       Just stuff -> Just stuff
+               -- 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
+
+       Nothing | isCheapUnfolding unf
+               -> is_con_app_maybe env (unfoldingTemplate unf)
+               where
+                 unf = idUnfolding v
+               -- However we do want to consult the unfolding as well,
+               -- for let-bound constructors!
+
+       other  -> Nothing
 
 is_con_app_maybe env (Lit lit)
   = Just (CV (LitAlt lit) [])