[commit: ghc] new-demand-to-merge: a problem with lambda-lifted join points solved (afcff01)

[Ilya Sergey]

Repository : ssh://darcs.haskell.org//srv/darcs/ghc

On branch  : new-demand-to-merge

http://hackage.haskell.org/trac/ghc/changeset/afcff010353289d639eba9fd0b5b8787c1a6e2c1

>---------------------------------------------------------------

commit afcff010353289d639eba9fd0b5b8787c1a6e2c1
Author: Ilya Sergey <Ilya.Sergey at cs.kuleuven.be>
Date:   Tue Sep 25 19:17:07 2012 +0100

    a problem with lambda-lifted join points solved

>---------------------------------------------------------------

 compiler/specialise/SpecConstr.lhs |    2 +-
 compiler/stranal/WwLib.lhs         |   12 +++++++++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/compiler/specialise/SpecConstr.lhs b/compiler/specialise/SpecConstr.lhs
index f842d3c..d7dec23 100644
--- a/compiler/specialise/SpecConstr.lhs
+++ b/compiler/specialise/SpecConstr.lhs
@@ -1404,7 +1404,7 @@ spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number)
 			     `setIdArity` count isId spec_lam_args
               spec_str   = calcSpecStrictness fn spec_lam_args pats
                 -- Conditionally use result of new worker-wrapper transform
-	      (spec_lam_args, spec_call_args) = mkWorkerArgs qvars body_ty
+	      (spec_lam_args, spec_call_args) = mkWorkerArgs qvars False body_ty
 	      	-- Usual w/w hack to avoid generating 
 	      	-- a spec_rhs of unlifted type and no args
 
diff --git a/compiler/stranal/WwLib.lhs b/compiler/stranal/WwLib.lhs
index 3590066..f471410 100644
--- a/compiler/stranal/WwLib.lhs
+++ b/compiler/stranal/WwLib.lhs
@@ -132,13 +132,14 @@ mkWwBodies :: DynFlags
 
 mkWwBodies dflags fun_ty demands res_info one_shots
   = do	{ let arg_info = demands `zip` (one_shots ++ repeat False)
+              all_one_shots = all snd arg_info
 	; (wrap_args, wrap_fn_args, work_fn_args, res_ty) <- mkWWargs emptyTvSubst fun_ty arg_info
 	; (work_args, wrap_fn_str,  work_fn_str) <- mkWWstr dflags wrap_args
 
         -- Do CPR w/w.  See Note [Always do CPR w/w]
 	; (wrap_fn_cpr, work_fn_cpr,  cpr_res_ty) <- mkWWcpr res_ty res_info
 
-	; let (work_lam_args, work_call_args) = mkWorkerArgs work_args cpr_res_ty
+	; let (work_lam_args, work_call_args) = mkWorkerArgs work_args all_one_shots cpr_res_ty
 	; return ([idDemandInfo v | v <- work_call_args, isId v],
                   wrap_fn_args . wrap_fn_cpr . wrap_fn_str . applyToVars work_call_args . Var,
                   mkLams work_lam_args. work_fn_str . work_fn_cpr . work_fn_args) }
@@ -183,14 +184,19 @@ We use the state-token type which generates no code.
 
 \begin{code}
 mkWorkerArgs :: [Var]
+             -> Bool    -- Whether all arguments are one-shot
 	     -> Type	-- Type of body
 	     -> ([Var],	-- Lambda bound args
 		 [Var])	-- Args at call site
-mkWorkerArgs args res_ty
+mkWorkerArgs args all_one_shot res_ty
     | any isId args || not (isUnLiftedType res_ty)
     = (args, args)
     | otherwise	
-    = (args ++ [voidArgId], args ++ [realWorldPrimId])
+    = (args ++ [newArg], args ++ [realWorldPrimId])
+    where
+      newArg = if all_one_shot 
+               then setOneShotLambda voidArgId
+               else voidArgId     
 \end{code}