Do not count void arguments when considering a function for loopification.

This fixes #11372 by omitting arguments with a void-type when checking
whether a self-recursive tail call can be optimized to a local jump.
Previously, a function taking a real argument and a State# token
would report an arity of 1 in the SelfLoopInfo in getCallMethod,
but a self-recursive call would apply it to 2 arguments, one of them
being the State# token, thus no local jump would be generated.
As the State# token is not represented by anything at runtime, we can
ignore it and thus trigger the loopification optimization.

Test Plan: ./validate

Reviewers: austin, bgamari, simonmar

Reviewed By: bgamari

Subscribers: simonmar, thomie

Differential Revision: https://phabricator.haskell.org/D1767

GHC Trac Issues: #11372

2 files changed, 46 insertions, 17 deletions

diff --git a/compiler/codeGen/StgCmmClosure.hs b/compiler/codeGen/StgCmmClosure.hs
index d3b9fac34a..97224c6660 100644
--- a/compiler/codeGen/StgCmmClosure.hs
+++ b/compiler/codeGen/StgCmmClosure.hs
@@ -496,6 +496,7 @@ getCallMethod :: DynFlags
                                 -- itself
               -> LambdaFormInfo -- Its info
               -> RepArity       -- Number of available arguments
+              -> RepArity       -- Number of them being void arguments
               -> CgLoc          -- Passed in from cgIdApp so that we can
                                 -- handle let-no-escape bindings and self-recursive
                                 -- tail calls using the same data constructor,
@@ -504,30 +505,34 @@ getCallMethod :: DynFlags
               -> Maybe SelfLoopInfo -- can we perform a self-recursive tail call?
               -> CallMethod
 
-getCallMethod dflags _ id _ n_args _cg_loc (Just (self_loop_id, block_id, args))
-  | gopt Opt_Loopification dflags, id == self_loop_id, n_args == length args
+getCallMethod dflags _ id _ n_args v_args _cg_loc
+              (Just (self_loop_id, block_id, args))
+  | gopt Opt_Loopification dflags
+  , id == self_loop_id
+  , n_args - v_args == length args
   -- If these patterns match then we know that:
   --   * loopification optimisation is turned on
   --   * function is performing a self-recursive call in a tail position
-  --   * number of parameters of the function matches functions arity.
-  -- See Note [Self-recursive tail calls] in StgCmmExpr for more details
+  --   * number of non-void parameters of the function matches functions arity.
+  -- See Note [Self-recursive tail calls] and Note [Void arguments in
+  -- self-recursive tail calls] in StgCmmExpr for more details
   = JumpToIt block_id args
 
-getCallMethod dflags name id (LFReEntrant _ arity _ _) n_args _cg_loc
+getCallMethod dflags name id (LFReEntrant _ arity _ _) n_args _v_args _cg_loc
               _self_loop_info
   | n_args == 0    = ASSERT( arity /= 0 )
                      ReturnIt        -- No args at all
   | n_args < arity = SlowCall        -- Not enough args
   | otherwise      = DirectEntry (enterIdLabel dflags name (idCafInfo id)) arity
 
-getCallMethod _ _name _ LFUnLifted n_args _cg_loc _self_loop_info
+getCallMethod _ _name _ LFUnLifted n_args _v_args _cg_loc _self_loop_info
   = ASSERT( n_args == 0 ) ReturnIt
 
-getCallMethod _ _name _ (LFCon _) n_args _cg_loc _self_loop_info
+getCallMethod _ _name _ (LFCon _) n_args _v_args _cg_loc _self_loop_info
   = ASSERT( n_args == 0 ) ReturnIt
 
 getCallMethod dflags name id (LFThunk _ _ updatable std_form_info is_fun)
-              n_args _cg_loc _self_loop_info
+              n_args _v_args _cg_loc _self_loop_info
   | is_fun      -- it *might* be a function, so we must "call" it (which is always safe)
   = SlowCall    -- We cannot just enter it [in eval/apply, the entry code
                 -- is the fast-entry code]
@@ -558,18 +563,18 @@ getCallMethod dflags name id (LFThunk _ _ updatable std_form_info is_fun)
     DirectEntry (thunkEntryLabel dflags name (idCafInfo id) std_form_info
                 updatable) 0
 
-getCallMethod _ _name _ (LFUnknown True) _n_arg _cg_locs _self_loop_info
+getCallMethod _ _name _ (LFUnknown True) _n_arg _v_args _cg_locs _self_loop_info
   = SlowCall -- might be a function
 
-getCallMethod _ name _ (LFUnknown False) n_args _cg_loc _self_loop_info
+getCallMethod _ name _ (LFUnknown False) n_args _v_args _cg_loc _self_loop_info
   = ASSERT2( n_args == 0, ppr name <+> ppr n_args )
     EnterIt -- Not a function
 
-getCallMethod _ _name _ LFLetNoEscape _n_args (LneLoc blk_id lne_regs)
+getCallMethod _ _name _ LFLetNoEscape _n_args _v_args (LneLoc blk_id lne_regs)
               _self_loop_info
   = JumpToIt blk_id lne_regs
 
-getCallMethod _ _ _ _ _ _ _ = panic "Unknown call method"
+getCallMethod _ _ _ _ _ _ _ _ = panic "Unknown call method"
 
 -----------------------------------------------------------------------------
 --                staticClosureRequired
diff --git a/compiler/codeGen/StgCmmExpr.hs b/compiler/codeGen/StgCmmExpr.hs
index c4ff11a1d0..8e74c22fba 100644
--- a/compiler/codeGen/StgCmmExpr.hs
+++ b/compiler/codeGen/StgCmmExpr.hs
@@ -697,8 +697,10 @@ cgIdApp fun_id args = do
         fun_name    = idName    cg_fun_id
         fun         = idInfoToAmode fun_info
         lf_info     = cg_lf         fun_info
+        n_args      = length args
+        v_args      = length $ filter (isVoidTy . stgArgType) args
         node_points dflags = nodeMustPointToIt dflags lf_info
-    case (getCallMethod dflags fun_name cg_fun_id lf_info (length args) (cg_loc fun_info) self_loop_info) of
+    case getCallMethod dflags fun_name cg_fun_id lf_info n_args v_args (cg_loc fun_info) self_loop_info of
 
             -- A value in WHNF, so we can just return it.
         ReturnIt -> emitReturn [fun]    -- ToDo: does ReturnIt guarantee tagged?
@@ -802,14 +804,36 @@ cgIdApp fun_id args = do
 --     of call will be generated. getCallMethod decides to generate a self
 --     recursive tail call when (a) environment stores information about
 --     possible self tail-call; (b) that tail call is to a function currently
---     being compiled; (c) number of passed arguments is equal to function's
---     arity. (d) loopification is turned on via -floopification command-line
---     option.
+--     being compiled; (c) number of passed non-void arguments is equal to
+--     function's arity. (d) loopification is turned on via -floopification
+--     command-line option.
 --
 --   * Command line option to turn loopification on and off is implemented in
 --     DynFlags.
 --
-
+--
+-- Note [Void arguments in self-recursive tail calls]
+-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+--
+-- State# tokens can get in the way of the loopification optimization as seen in
+-- #11372. Consider this:
+--
+-- foo :: [a]
+--     -> (a -> State# s -> (# State s, Bool #))
+--     -> State# s
+--     -> (# State# s, Maybe a #)
+-- foo [] f s = (# s, Nothing #)
+-- foo (x:xs) f s = case f x s of
+--      (# s', b #) -> case b of
+--          True -> (# s', Just x #)
+--          False -> foo xs f s'
+--
+-- We would like to compile the call to foo as a local jump instead of a call
+-- (see Note [Self-recursive tail calls]). However, the generated function has
+-- an arity of 2 while we apply it to 3 arguments, one of them being of void
+-- type. Thus, we mustn't count arguments of void type when checking whether
+-- we can turn a call into a self-recursive jump.
+--
 
 emitEnter :: CmmExpr -> FCode ReturnKind
 emitEnter fun = do