-lvlExpr ctxt_lvl env (_, AnnNote note expr)
- = lvlExpr ctxt_lvl env expr `thenLvl` \ expr' ->
- returnLvl (Note note expr')
-
-lvlExpr ctxt_lvl env (_, AnnCast expr co)
- = lvlExpr ctxt_lvl env expr `thenLvl` \ expr' ->
- returnLvl (Cast expr' co)
+lvlExpr ctxt_lvl env expr@(_, AnnApp _ _) = do
+ let
+ (fun, args) = collectAnnArgs expr
+ --
+ case fun of
+ -- float out partial applications. This is very beneficial
+ -- in some cases (-7% runtime -4% alloc over nofib -O2).
+ -- In order to float a PAP, there must be a function at the
+ -- head of the application, and the application must be
+ -- over-saturated with respect to the function's arity.
+ (_, AnnVar f) | floatPAPs env &&
+ arity > 0 && arity < n_val_args ->
+ do
+ let (lapp, rargs) = left (n_val_args - arity) expr []
+ rargs' <- mapM (lvlMFE False ctxt_lvl env) rargs
+ lapp' <- lvlMFE False ctxt_lvl env lapp
+ return (foldl App lapp' rargs')
+ where
+ n_val_args = count (isValArg . deAnnotate) args
+ arity = idArity f
+
+ -- separate out the PAP that we are floating from the extra
+ -- arguments, by traversing the spine until we have collected
+ -- (n_val_args - arity) value arguments.
+ left 0 e rargs = (e, rargs)
+ left n (_, AnnApp f a) rargs
+ | isValArg (deAnnotate a) = left (n-1) f (a:rargs)
+ | otherwise = left n f (a:rargs)
+ left _ _ _ = panic "SetLevels.lvlExpr.left"
+
+ -- No PAPs that we can float: just carry on with the
+ -- arguments and the function.
+ _otherwise -> do
+ args' <- mapM (lvlMFE False ctxt_lvl env) args
+ fun' <- lvlExpr ctxt_lvl env fun
+ return (foldl App fun' args')
+
+lvlExpr ctxt_lvl env (_, AnnNote note expr) = do
+ expr' <- lvlExpr ctxt_lvl env expr
+ return (Note note expr')
+
+lvlExpr ctxt_lvl env (_, AnnCast expr (_, co)) = do
+ expr' <- lvlExpr ctxt_lvl env expr
+ return (Cast expr' co)