Make the simplifier propagate strictness through casts

E.g.  (f e1 |> g) e2
If f is strict in two aguments, we want to see that in e2

Hence ArgSpec in SimplUtils

2 files changed, 83 insertions, 41 deletions

diff --git a/compiler/simplCore/SimplUtils.lhs b/compiler/simplCore/SimplUtils.lhs
index c483f05494..92874de4a3 100644
--- a/compiler/simplCore/SimplUtils.lhs
+++ b/compiler/simplCore/SimplUtils.lhs
@@ -15,15 +15,17 @@ module SimplUtils (
         simplEnvForGHCi, updModeForInlineRules,
 
         -- The continuation type
-        SimplCont(..), DupFlag(..), ArgInfo(..),
+        SimplCont(..), DupFlag(..), 
         isSimplified,
         contIsDupable, contResultType, contInputType,
         contIsTrivial, contArgs, dropArgs,
-        pushSimplifiedArgs, countValArgs, countArgs, addArgTo,
+        pushSimplifiedArgs, countValArgs, countArgs, 
         mkBoringStop, mkRhsStop, mkLazyArgStop, contIsRhsOrArg,
-        interestingCallContext,
+        interestingCallContext, interestingArg, 
 
-        interestingArg, mkArgInfo,
+        -- ArgInfo
+        ArgInfo(..), ArgSpec(..), mkArgInfo, addArgTo, addCastTo, 
+        argInfoExpr, argInfoValArgs,
 
         abstractFloats
     ) where
@@ -132,7 +134,7 @@ data SimplCont
 data ArgInfo
   = ArgInfo {
         ai_fun   :: OutId,      -- The function
-        ai_args  :: [OutExpr],  -- ...applied to these args (which are in *reverse* order)
+        ai_args  :: [ArgSpec],  -- ...applied to these args (which are in *reverse* order)
         ai_type  :: OutType,    -- Type of (f a1 ... an)
 
         ai_rules :: [CoreRule], -- Rules for this function
@@ -149,10 +151,38 @@ data ArgInfo
                                 --   Always infinite
     }
 
+data ArgSpec = ValArg OutExpr       -- Apply to this
+             | CastBy OutCoercion   -- Cast by this
+
+instance Outputable ArgSpec where
+  ppr (ValArg e) = ptext (sLit "ValArg") <+> ppr e
+  ppr (CastBy c) = ptext (sLit "CastBy") <+> ppr c
+
 addArgTo :: ArgInfo -> OutExpr -> ArgInfo
-addArgTo ai arg = ai { ai_args = arg : ai_args ai
+addArgTo ai arg = ai { ai_args = ValArg arg : ai_args ai
                      , ai_type = applyTypeToArg (ai_type ai) arg  }
 
+addCastTo :: ArgInfo -> OutCoercion -> ArgInfo
+addCastTo ai co = ai { ai_args = CastBy co : ai_args ai
+                     , ai_type = pSnd (coercionKind co) }
+
+argInfoValArgs :: SimplEnv -> [ArgSpec] -> SimplCont -> ([OutExpr], SimplCont)
+argInfoValArgs env args cont
+  = go args [] cont
+  where
+    go :: [ArgSpec] -> [OutExpr] -> SimplCont -> ([OutExpr], SimplCont)
+    go (ValArg e  : as) acc cont = go as (e:acc) cont
+    go (CastBy co : as) acc cont = go as [] (CoerceIt co (pushSimplifiedArgs env acc cont))
+    go []               acc cont = (acc, cont)
+
+argInfoExpr :: OutId -> [ArgSpec] -> OutExpr
+argInfoExpr fun args
+  = go args
+  where
+    go []               = Var fun
+    go (ValArg a : as)  = go as `App` a
+    go (CastBy co : as) = mkCast (go as) co
+
 instance Outputable SimplCont where
   ppr (Stop ty interesting)          = ptext (sLit "Stop") <> brackets (ppr interesting) <+> ppr ty
   ppr (ApplyTo dup arg _ cont)       = ((ptext (sLit "ApplyTo") <+> ppr dup <+> pprParendExpr arg)
@@ -258,21 +288,27 @@ countArgs (ApplyTo _ _ _ cont) = 1 + countArgs cont
 countArgs _                    = 0
 
 contArgs :: SimplCont -> (Bool, [ArgSummary], SimplCont)
--- Uses substitution to turn each arg into an OutExpr
-contArgs cont@(ApplyTo {})
-  = case go [] cont of { (args, cont') -> (False, args, cont') }
+-- Summarises value args, discards type args and coercions
+-- The returned continuation of the call is only used to 
+-- answer questions like "are you interesting?"
+contArgs cont
+  | lone cont = (True, [], cont)
+  | otherwise = go [] cont
   where
+    lone (ApplyTo {})  = False  -- See Note [Lone variables] in CoreUnfold
+    lone (CoerceIt {}) = False
+    lone _             = True
+
     go args (ApplyTo _ arg se cont)
-      | isTypeArg arg = go args                           cont
-      | otherwise     = go (is_interesting arg se : args) cont
-    go args cont      = (reverse args, cont)
+      | isTypeArg arg         = go args                           cont
+      | otherwise             = go (is_interesting arg se : args) cont
+    go args (CoerceIt _ cont) = go args cont
+    go args cont              = (False, reverse args, cont)
 
     is_interesting arg se = interestingArg (substExpr (text "contArgs") se arg)
                    -- Do *not* use short-cutting substitution here
                    -- because we want to get as much IdInfo as possible
 
-contArgs cont = (True, [], cont)
-
 pushSimplifiedArgs :: SimplEnv -> [CoreExpr] -> SimplCont -> SimplCont
 pushSimplifiedArgs _env []         cont = cont
 pushSimplifiedArgs env  (arg:args) cont = ApplyTo Simplified arg env (pushSimplifiedArgs env args cont)
diff --git a/compiler/simplCore/Simplify.lhs b/compiler/simplCore/Simplify.lhs
index 0bc05f3985..f0f894d744 100644
--- a/compiler/simplCore/Simplify.lhs
+++ b/compiler/simplCore/Simplify.lhs
@@ -33,7 +33,6 @@ import CoreUtils
 import qualified CoreSubst
 import CoreArity
 import Rules            ( lookupRule, getRules )
-import BasicTypes       ( Arity )
 import TysPrim          ( realWorldStatePrimTy )
 import BasicTypes       ( TopLevelFlag(..), isTopLevel, RecFlag(..) )
 import MonadUtils       ( foldlM, mapAccumLM, liftIO )
@@ -537,6 +536,11 @@ These strange casts can happen as a result of case-of-case
 
 
 \begin{code}
+makeTrivialArg :: SimplEnv -> ArgSpec -> SimplM (SimplEnv, ArgSpec)
+makeTrivialArg env (ValArg e)  = do { (env', e') <- makeTrivial NotTopLevel env e
+                                    ; return (env', ValArg e') }
+makeTrivialArg env (CastBy co) = return (env, CastBy co)
+
 makeTrivial :: TopLevelFlag -> SimplEnv -> OutExpr -> SimplM (SimplEnv, OutExpr)
 -- Binds the expression to a variable, if it's not trivial, returning the variable
 makeTrivial top_lvl env expr = makeTrivialWithInfo top_lvl env vanillaIdInfo expr
@@ -1394,12 +1398,6 @@ completeCall env var cont
   = do  {   ------------- Try inlining ----------------
           dflags <- getDynFlags
         ; let  (lone_variable, arg_infos, call_cont) = contArgs cont
-                -- The args are OutExprs, obtained by *lazily* substituting
-                -- in the args found in cont.  These args are only examined
-                -- to limited depth (unless a rule fires).  But we must do
-                -- the substitution; rule matching on un-simplified args would
-                -- be bogus
-
                n_val_args = length arg_infos
                interesting_cont = interestingCallContext call_cont
                unfolding    = activeUnfolding env var
@@ -1448,9 +1446,12 @@ rebuildCall env (ArgInfo { ai_fun = fun, ai_args = rev_args, ai_strs = [] }) con
   | not (contIsTrivial cont)     -- Only do this if there is a non-trivial
   = return (env, castBottomExpr res cont_ty)  -- contination to discard, else we do it
   where                                       -- again and again!
-    res     = mkApps (Var fun) (reverse rev_args)
+    res     = argInfoExpr fun rev_args
     cont_ty = contResultType cont
 
+rebuildCall env info (CoerceIt co cont) 
+  = rebuildCall env (addCastTo info co) cont
+
 rebuildCall env info (ApplyTo dup_flag (Type arg_ty) se cont)
   = do { arg_ty' <- if isSimplified dup_flag then return arg_ty
                     else simplType (se `setInScope` env) arg_ty
@@ -1482,17 +1483,21 @@ rebuildCall env info@(ArgInfo { ai_encl = encl_rules, ai_type = fun_ty
         | otherwise              = BoringCtxt          -- Nothing interesting
 
 rebuildCall env (ArgInfo { ai_fun = fun, ai_args = rev_args, ai_rules = rules }) cont
+  | null rules
+  = rebuild env (argInfoExpr fun rev_args) cont      -- No rules, common case
+
+  | otherwise
   = do {  -- We've accumulated a simplified call in <fun,rev_args>
           -- so try rewrite rules; see Note [RULEs apply to simplified arguments]
           -- See also Note [Rules for recursive functions]
-        ; let args = reverse rev_args
-              env' = zapSubstEnv env
-        ; mb_rule <- tryRules env rules fun args cont
+        ; let env' = zapSubstEnv env
+              (args, cont') = argInfoValArgs env' rev_args cont
+        ; mb_rule <- tryRules env' rules fun args cont'
         ; case mb_rule of {
-             Just (n_args, rule_rhs) -> simplExprF env' rule_rhs $
-                                        pushSimplifiedArgs env' (drop n_args args) cont ;
-                 -- n_args says how many args the rule consumed
-           ; Nothing -> rebuild env (mkApps (Var fun) args) cont      -- No rules
+             Just (rule_rhs, cont'') -> simplExprF env' rule_rhs cont''
+
+                 -- Rules don't match
+           ; Nothing -> rebuild env (argInfoExpr fun rev_args) cont      -- No rules
     } }
 \end{code}
 
@@ -1552,8 +1557,9 @@ all this at once is TOO HARD!
 \begin{code}
 tryRules :: SimplEnv -> [CoreRule]
          -> Id -> [OutExpr] -> SimplCont
-         -> SimplM (Maybe (Arity, CoreExpr)) -- The arity is the number of
-                                             -- args consumed by the rule
+         -> SimplM (Maybe (CoreExpr, SimplCont)) 
+-- The SimplEnv already has zapSubstEnv applied to it
+
 tryRules env rules fn args call_cont
   | null rules
   = return Nothing
@@ -1563,11 +1569,13 @@ tryRules env rules fn args call_cont
                          fn args rules of {
            Nothing               -> return Nothing ;   -- No rule matches
            Just (rule, rule_rhs) ->
-
              do { checkedTick (RuleFired (ru_name rule))
-                ; dflags <- getDynFlags
                 ; dump dflags rule rule_rhs
-                ; return (Just (ruleArity rule, rule_rhs)) }}}
+                ; let cont' = pushSimplifiedArgs env
+                                                 (drop (ruleArity rule) args) 
+                                                 call_cont
+                      -- (ruleArity rule) says how many args the rule consumed
+                ; return (Just (rule_rhs, cont')) }}}
   where
     dump dflags rule rule_rhs
       | dopt Opt_D_dump_rule_rewrites dflags
@@ -1586,7 +1594,6 @@ tryRules env rules fn args call_cont
 
     log_rule dflags flag hdr details = liftIO . dumpSDoc dflags flag "" $
       sep [text hdr, nest 4 details]
-
 \end{code}
 
 Note [Rules for recursive functions]
@@ -1858,17 +1865,16 @@ rebuildCase env scrut case_bndr [(_, bndrs, rhs)] cont
 rebuildCase env scrut case_bndr alts@[(_, bndrs, rhs)] cont
   | all isDeadBinder (case_bndr : bndrs)  -- So this is just 'seq'
   = do { let rhs' = substExpr (text "rebuild-case") env rhs
+             env' = zapSubstEnv env
              out_args = [Type (substTy env (idType case_bndr)),
                          Type (exprType rhs'), scrut, rhs']
                       -- Lazily evaluated, so we don't do most of this
 
        ; rule_base <- getSimplRules
-       ; mb_rule <- tryRules env (getRules rule_base seqId) seqId out_args cont
+       ; mb_rule <- tryRules env' (getRules rule_base seqId) seqId out_args cont
        ; case mb_rule of
-           Just (n_args, res) -> simplExprF (zapSubstEnv env)
-                                            (mkApps res (drop n_args out_args))
-                                            cont
-           Nothing -> reallyRebuildCase env scrut case_bndr alts cont }
+           Just (rule_rhs, cont') -> simplExprF env' rule_rhs cont'
+           Nothing                -> reallyRebuildCase env scrut case_bndr alts cont }
 
 rebuildCase env scrut case_bndr alts cont
   = reallyRebuildCase env scrut case_bndr alts cont
@@ -2315,7 +2321,7 @@ mkDupableCont env cont@(StrictBind {})
 mkDupableCont env (StrictArg info cci cont)
         -- See Note [Duplicating StrictArg]
   = do { (env', dup, nodup) <- mkDupableCont env cont
-       ; (env'', args')     <- mapAccumLM (makeTrivial NotTopLevel) env' (ai_args info)
+       ; (env'', args')     <- mapAccumLM makeTrivialArg env' (ai_args info)
        ; return (env'', StrictArg (info { ai_args = args' }) cci dup, nodup) }
 
 mkDupableCont env (ApplyTo _ arg se cont)