Fix a bug where finaliseArgBoxities wasn't looking through casts.wip/andreask/opaque-boxity-fix

If the return type of a function was a newtype then we would fail to
adjust the demands resulting in a panic in W/W when looking at opaque
bindings.

Fixes #22502

5 files changed, 35 insertions, 8 deletions

diff --git a/compiler/GHC/Core.hs b/compiler/GHC/Core.hs
index 92b34ffc21..78cc26063a 100644
--- a/compiler/GHC/Core.hs
+++ b/compiler/GHC/Core.hs
@@ -41,6 +41,7 @@ module GHC.Core (
         bindersOf, bindersOfBinds, rhssOfBind, rhssOfAlts,
         foldBindersOfBindStrict, foldBindersOfBindsStrict,
         collectBinders, collectTyBinders, collectTyAndValBinders,
+        collectBindersThroughCasts,
         collectNBinders, collectNValBinders_maybe,
         collectArgs, stripNArgs, collectArgsTicks, flattenBinds,
         collectFunSimple,
@@ -1954,6 +1955,7 @@ flattenBinds []                   = []
 collectBinders         :: Expr b   -> ([b],     Expr b)
 collectTyBinders       :: CoreExpr -> ([TyVar], CoreExpr)
 collectValBinders      :: CoreExpr -> ([Id],    CoreExpr)
+collectBindersThroughCasts :: Expr Var -> ([Var], Expr Var)
 collectTyAndValBinders :: CoreExpr -> ([TyVar], [Id], CoreExpr)
 
 -- | Strip off exactly N leading lambdas (type or value).
@@ -1979,6 +1981,14 @@ collectValBinders expr
     go ids (Lam b e) | isId b = go (b:ids) e
     go ids body               = (reverse ids, body)
 
+-- | Look through casts when collecting binders
+collectBindersThroughCasts expr
+  = go [] expr
+  where
+    go ids (Cast e _) = go ids e
+    go ids (Lam b e)  = go (b:ids) e
+    go ids body       = (reverse ids, body)
+
 collectTyAndValBinders expr
   = (tvs, ids, body)
   where
diff --git a/compiler/GHC/Core/Opt/DmdAnal.hs b/compiler/GHC/Core/Opt/DmdAnal.hs
index 3738f8b3ed..d14c73628c 100644
--- a/compiler/GHC/Core/Opt/DmdAnal.hs
+++ b/compiler/GHC/Core/Opt/DmdAnal.hs
@@ -48,7 +48,6 @@ import GHC.Types.Unique.Set
 import GHC.Types.Unique.MemoFun
 import GHC.Types.RepType
 
-
 {-
 ************************************************************************
 *                                                                      *
@@ -1922,7 +1921,7 @@ finaliseArgBoxities env fn arity rhs div
     -- uses the info on the binders directly.
   where
     opts            = ae_opts env
-    (bndrs, _body)  = collectBinders rhs
+    (bndrs, _body)  = collectBindersThroughCasts rhs
     unarise_arity   = sum [ unariseArity (idType b) | b <- bndrs, isId b ]
     max_wkr_args    = dmd_max_worker_args opts `max` unarise_arity
                       -- This is the budget initialisation step of
@@ -1947,17 +1946,19 @@ finaliseArgBoxities env fn arity rhs div
         --     catch newtype dictionaries too.
         -- NB: even for bottoming functions, don't unbox dictionaries
 
-      | is_bot_fn = unboxDeeplyDmd dmd
-        -- See Note [Boxity for bottoming functions], case (B)
-
       | is_opaque = trimBoxity dmd
         -- See Note [OPAQUE pragma]
         -- See Note [The OPAQUE pragma and avoiding the reboxing of arguments]
 
+      | is_bot_fn = unboxDeeplyDmd dmd
+        -- See Note [Boxity for bottoming functions], case (B)
+
+
       | otherwise = dmd
       where
         dmd       = idDemandInfo bndr
-        is_opaque = isOpaquePragma (idInlinePragma fn)
+
+    is_opaque = isOpaquePragma (idInlinePragma fn)
 
     -- is_bot_fn:  see Note [Boxity for bottoming functions]
     is_bot_fn = div == botDiv
@@ -2017,6 +2018,7 @@ finaliseArgBoxities env fn arity rhs div
     add_demands :: [Demand] -> CoreExpr -> CoreExpr
     -- Attach the demands to the outer lambdas of this expression
     add_demands [] e = e
+    add_demands (dmds) (Cast e co) = Cast (add_demands dmds e) co
     add_demands (dmd:dmds) (Lam v e)
       | isTyVar v = Lam v (add_demands (dmd:dmds) e)
       | otherwise = Lam (v `setIdDemandInfo` dmd) (add_demands dmds e)
diff --git a/compiler/GHC/Utils/Misc.hs b/compiler/GHC/Utils/Misc.hs
index a115c61336..50ede55909 100644
--- a/compiler/GHC/Utils/Misc.hs
+++ b/compiler/GHC/Utils/Misc.hs
@@ -637,8 +637,9 @@ all2 _ []     []     = True
 all2 p (x:xs) (y:ys) = p x y && all2 p xs ys
 all2 _ _      _      = False
 
--- Count the number of times a predicate is true
-
+-- | Count the number of times a predicate is true
+--
+-- A manually fused alternative to @length . filter p@
 count :: (a -> Bool) -> [a] -> Int
 count p = go 0
   where go !n [] = n
diff --git a/testsuite/tests/simplCore/should_compile/OpaqueNoWW2.hs b/testsuite/tests/simplCore/should_compile/OpaqueNoWW2.hs
new file mode 100644
index 0000000000..92f965d98b
--- /dev/null
+++ b/testsuite/tests/simplCore/should_compile/OpaqueNoWW2.hs
@@ -0,0 +1,13 @@
+module OpaqueNoWW2 where
+
+{-# LANGUAGE MagicHash #-}
+module M where
+
+import GHC.Exts
+import GHC.IO
+
+data T a = MkT !Bool !a
+
+fun :: T a -> IO a
+{-# OPAQUE fun #-}
+fun (MkT _ x) = IO $ \s -> noinline seq# x s
diff --git a/testsuite/tests/simplCore/should_compile/all.T b/testsuite/tests/simplCore/should_compile/all.T
index c5f63d6e7a..20abf0dab6 100644
--- a/testsuite/tests/simplCore/should_compile/all.T
+++ b/testsuite/tests/simplCore/should_compile/all.T
@@ -398,6 +398,7 @@ test('OpaqueNoSpecConstr', [ req_interp, grep_errmsg(r'$sloop') ], compile, ['-O
 test('OpaqueNoSpecialise', [ grep_errmsg(r'$sf') ], compile, ['-O -ddump-simpl -dsuppress-uniques'])
 test('OpaqueNoStrictArgWW', [ grep_errmsg(r'$wf') ], compile, ['-O -fworker-wrapper-cbv -ddump-simpl -dsuppress-uniques'])
 test('OpaqueNoWW', [ grep_errmsg(r'$wf') ], compile, ['-O -ddump-simpl -dsuppress-uniques'])
+test('OpaqueNoWW2', [ grep_errmsg(r'$wf') ], compile, ['-O -ddump-simpl -dsuppress-uniques'])
 
 test('T21144',  normal, compile, ['-O'])