Allow the demand analyser to unpack tuple and equality dictionarieswip/T23398

Addresses #23398. The demand analyser usually does not unpack class
dictionaries: see Note [Do not unbox class dictionaries] in
GHC.Core.Opt.DmdAnal.

This patch makes an exception for tuple dictionaries and equality
dictionaries, for reasons explained in wrinkles (DNB1) and (DNB2) of
the above Note.

5 files changed, 241 insertions, 69 deletions

diff --git a/compiler/GHC/Core/Opt/DmdAnal.hs b/compiler/GHC/Core/Opt/DmdAnal.hs
index 0b74a9e1d2..f308d6c0e3 100644
--- a/compiler/GHC/Core/Opt/DmdAnal.hs
+++ b/compiler/GHC/Core/Opt/DmdAnal.hs
@@ -16,37 +16,41 @@ where
 
 import GHC.Prelude
 
-import GHC.Core.Opt.WorkWrap.Utils
 import GHC.Types.Demand   -- All of it
+
 import GHC.Core
-import GHC.Core.Multiplicity ( scaledThing )
-import GHC.Utils.Outputable
-import GHC.Types.Var.Env
-import GHC.Types.Var.Set
-import GHC.Types.Basic
-import Data.List        ( mapAccumL )
 import GHC.Core.DataCon
-import GHC.Types.ForeignCall ( isSafeForeignCall )
-import GHC.Types.Id
 import GHC.Core.Utils
 import GHC.Core.TyCon
 import GHC.Core.Type
-import GHC.Core.Predicate( isClassPred )
+import GHC.Core.Predicate( isEqualityClass, isCTupleClass )
 import GHC.Core.FVs      ( rulesRhsFreeIds, bndrRuleAndUnfoldingIds )
 import GHC.Core.Coercion ( Coercion )
 import GHC.Core.TyCo.FVs     ( coVarsOfCos )
 import GHC.Core.TyCo.Compare ( eqType )
+import GHC.Core.Multiplicity ( scaledThing )
 import GHC.Core.FamInstEnv
 import GHC.Core.Opt.Arity ( typeArity )
-import GHC.Utils.Misc
-import GHC.Utils.Panic
-import GHC.Utils.Panic.Plain
+import GHC.Core.Opt.WorkWrap.Utils
+
 import GHC.Builtin.PrimOps
 import GHC.Builtin.Types.Prim ( realWorldStatePrimTy )
+
 import GHC.Types.Unique.Set
 import GHC.Types.Unique.MemoFun
 import GHC.Types.RepType
+import GHC.Types.ForeignCall ( isSafeForeignCall )
+import GHC.Types.Id
+import GHC.Types.Var.Env
+import GHC.Types.Var.Set
+import GHC.Types.Basic
 
+import GHC.Utils.Misc
+import GHC.Utils.Panic
+import GHC.Utils.Panic.Plain
+import GHC.Utils.Outputable
+
+import Data.List        ( mapAccumL )
 
 {-
 ************************************************************************
@@ -1499,7 +1503,7 @@ bounds-checking.
 
 So we want to give `indexError` a signature like `<1!P(!S,!S)><1!S><S!S>b`
 where the !S (meaning Poly Unboxed C1N) says that the polymorphic arguments
-are unboxed (recursively).  The wrapper for `indexError` won't /acutally/
+are unboxed (recursively).  The wrapper for `indexError` won't /actually/
 unbox them (because their polymorphic type doesn't allow that) but when
 demand-analysing /callers/, we'll behave as if that call needs the args
 unboxed.
@@ -1782,39 +1786,6 @@ applying the strictness demands to the final result of DmdAnal. The result is
 that we get the strict demand signature we wanted even if we can't float
 the case on `x` up through the case on `burble`.
 
-Note [Do not unbox class dictionaries]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-We never unbox class dictionaries in worker/wrapper.
-
-1. INLINABLE functions
-   If we have
-      f :: Ord a => [a] -> Int -> a
-      {-# INLINABLE f #-}
-   and we worker/wrapper f, we'll get a worker with an INLINABLE pragma
-   (see Note [Worker/wrapper for INLINABLE functions] in GHC.Core.Opt.WorkWrap),
-   which can still be specialised by the type-class specialiser, something like
-      fw :: Ord a => [a] -> Int# -> a
-
-   BUT if f is strict in the Ord dictionary, we might unpack it, to get
-      fw :: (a->a->Bool) -> [a] -> Int# -> a
-   and the type-class specialiser can't specialise that. An example is #6056.
-
-   Historical note: #14955 describes how I got this fix wrong the first time.
-   I got aware of the issue in T5075 by the change in boxity of loop between
-   demand analysis runs.
-
-2. -fspecialise-aggressively.  As #21286 shows, the same phenomenon can occur
-   occur without INLINABLE, when we use -fexpose-all-unfoldings and
-   -fspecialise-aggressively to do vigorous cross-module specialisation.
-
-3. #18421 found that unboxing a dictionary can also make the worker less likely
-   to inline; the inlining heuristics seem to prefer to inline a function
-   applied to a dictionary over a function applied to a bunch of functions.
-
-TL;DR we /never/ unbox class dictionaries. Unboxing the dictionary, and passing
-a raft of higher-order functions isn't a huge win anyway -- you really want to
-specialise the function.
-
 Note [Worker argument budget]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 In 'finaliseArgBoxities' we don't want to generate workers with zillions of
@@ -1998,22 +1969,13 @@ finaliseArgBoxities env fn threshold_arity rhs_dmds div rhs
 
     arg_triples :: [(Type, StrictnessMark, Demand)]
     arg_triples = take threshold_arity $
-                  [ (bndr_ty, NotMarkedStrict, get_dmd bndr bndr_ty)
-                  | bndr <- bndrs
-                  , isRuntimeVar bndr, let bndr_ty = idType bndr ]
-
-    get_dmd :: Id -> Type -> Demand
-    get_dmd bndr bndr_ty
-      | isClassPred bndr_ty = trimBoxity dmd
-        -- See Note [Do not unbox class dictionaries]
-        -- NB: 'ty' has not been normalised, so this will (rightly)
-        --     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)
-
-      | otherwise = dmd
+                  [ (idType bndr, NotMarkedStrict, get_dmd bndr)
+                  | bndr <- bndrs, isRuntimeVar bndr ]
+
+    get_dmd :: Id -> Demand
+    get_dmd bndr
+      | is_bot_fn = unboxDeeplyDmd dmd -- See Note [Boxity for bottoming functions],
+      | otherwise = dmd                --     case (B)
       where
         dmd = idDemandInfo bndr
 
@@ -2119,6 +2081,12 @@ wantToUnboxArg env ty str_mark dmd@(n :* _)
          -- isMarkedStrict: see Note [Unboxing evaluated arguments] in DmdAnal
        -> DontUnbox
 
+       | doNotUnbox ty
+       -> DontUnbox  -- See Note [Do not unbox class dictionaries]
+                     -- NB: 'ty' has not been normalised, so this will (rightly)
+                     --     catch newtype dictionaries too.
+                     -- NB: even for bottoming functions, don't unbox dictionaries
+
        | DefinitelyRecursive <- ae_rec_dc env dc
          -- See Note [Which types are unboxed?]
          -- and Note [Demand analysis for recursive data constructors]
@@ -2129,6 +2097,76 @@ wantToUnboxArg env ty str_mark dmd@(n :* _)
                         (dataConRepStrictness dc)
                         dmds)
 
+
+doNotUnbox :: Type -> Bool
+-- Do not unbox class dictionaries, except equality classes and tuples
+-- Note [Do not unbox class dictionaries]
+doNotUnbox arg_ty
+  = case tyConAppTyCon_maybe arg_ty of
+      Just tc | Just cls <- tyConClass_maybe tc
+              -> not (isEqualityClass cls || isCTupleClass cls)
+       -- See (DNB2) and (DNB1) in Note [Do not unbox class dictionaries]
+
+      _ -> False
+
+{- Note [Do not unbox class dictionaries]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+We never unbox class dictionaries in worker/wrapper.
+
+1. INLINABLE functions
+   If we have
+      f :: Ord a => [a] -> Int -> a
+      {-# INLINABLE f #-}
+   and we worker/wrapper f, we'll get a worker with an INLINABLE pragma
+   (see Note [Worker/wrapper for INLINABLE functions] in GHC.Core.Opt.WorkWrap),
+   which can still be specialised by the type-class specialiser, something like
+      fw :: Ord a => [a] -> Int# -> a
+
+   BUT if f is strict in the Ord dictionary, we might unpack it, to get
+      fw :: (a->a->Bool) -> [a] -> Int# -> a
+   and the type-class specialiser can't specialise that. An example is #6056.
+
+   Historical note: #14955 describes how I got this fix wrong the first time.
+   I got aware of the issue in T5075 by the change in boxity of loop between
+   demand analysis runs.
+
+2. -fspecialise-aggressively.  As #21286 shows, the same phenomenon can occur
+   occur without INLINABLE, when we use -fexpose-all-unfoldings and
+   -fspecialise-aggressively to do vigorous cross-module specialisation.
+
+3. #18421 found that unboxing a dictionary can also make the worker less likely
+   to inline; the inlining heuristics seem to prefer to inline a function
+   applied to a dictionary over a function applied to a bunch of functions.
+
+TL;DR we /never/ unbox class dictionaries. Unboxing the dictionary, and passing
+a raft of higher-order functions isn't a huge win anyway -- you really want to
+specialise the function.
+
+Wrinkle (DNB1): we /do/ want to unbox tuple dictionaries (#23398)
+     f :: (% Eq a, Show a %) => blah
+  with -fdicts-strict it is great to unbox to
+     $wf :: Eq a => Show a => blah
+  (where I have written out the currying explicitly).  Now we can specialise
+  $wf on the Eq or Show dictionary.  Nothing is lost.
+
+  And something is gained.  It is possible that `f` will look like this:
+     f = /\a. \d:(% Eq a, Show a %). ... f @a (% sel1 d, sel2 d %)...
+  where there is a recurive call to `f`, or to another function that takes the
+  same tuple dictionary, but where the tuple is built from the components of
+  `d`.  The Simplier does not fix this.  But if we unpacked the dictionary
+  we'd get
+     $wf = /\a. \(d1:Eq a) (d2:Show a). let d = (% d1, d2 %)
+             in ...f @a (% sel1 d, sel2 d %)
+  and all the tuple building and taking apart will disappear.
+
+Wrinkle (DNB2): we /do/ wnat to unbox equality dictionaries,
+  for (~), (~~), and Coercible (#23398).  Their payload is a single unboxed
+  coercion.  We never want to specialise on `(t1 ~ t2)`.  All that would do is
+  to make a copy of the function's RHS with a particular coercion.  Unlike
+  normal class methods, that does not unlock any new optimisation
+  opportunities in the specialised RHS.
+-}
+
 {- *********************************************************************
 *                                                                      *
                       Fixpoints
diff --git a/compiler/GHC/Core/Predicate.hs b/compiler/GHC/Core/Predicate.hs
index 2fc07e1be1..d6d5dd6520 100644
--- a/compiler/GHC/Core/Predicate.hs
+++ b/compiler/GHC/Core/Predicate.hs
@@ -20,7 +20,7 @@ module GHC.Core.Predicate (
 
   -- Class predicates
   mkClassPred, isDictTy, typeDeterminesValue,
-  isClassPred, isEqPredClass, isCTupleClass,
+  isClassPred, isEqPredClass, isCTupleClass, isEqualityClass,
   getClassPredTys, getClassPredTys_maybe,
   classMethodTy, classMethodInstTy,
 
@@ -219,11 +219,6 @@ isEvVarType :: Type -> Bool
 -- See Note [Evidence for quantified constraints]
 isEvVarType ty = isCoVarType ty || isPredTy ty
 
-isEqPredClass :: Class -> Bool
--- True of (~) and (~~)
-isEqPredClass cls =  cls `hasKey` eqTyConKey
-                  || cls `hasKey` heqTyConKey
-
 isClassPred :: PredType -> Bool
 isClassPred ty = case tyConAppTyCon_maybe ty of
     Just tc -> isClassTyCon tc
@@ -245,6 +240,20 @@ isEqPrimPred ty = isCoVarType ty
 isCTupleClass :: Class -> Bool
 isCTupleClass cls = isTupleTyCon (classTyCon cls)
 
+isEqPredClass :: Class -> Bool
+-- True of (~) and (~~)
+isEqPredClass cls =  cls `hasKey` eqTyConKey
+                  || cls `hasKey` heqTyConKey
+
+isEqualityClass :: Class -> Bool
+-- True of (~), (~~), and Coercible
+-- These all have a single primitive-equality superclass, either (~N# or ~R#)
+isEqualityClass cls
+  = cls `hasKey` heqTyConKey
+    || cls `hasKey` eqTyConKey
+    || cls `hasKey` coercibleTyConKey
+
+
 {- *********************************************************************
 *                                                                      *
               Implicit parameters
diff --git a/testsuite/tests/stranal/should_compile/T23398.hs b/testsuite/tests/stranal/should_compile/T23398.hs
new file mode 100644
index 0000000000..2a952942b8
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T23398.hs
@@ -0,0 +1,15 @@
+{-# OPTIONS_GHC -fdicts-strict #-}
+module T23398 where
+
+type PairDict a = (Eq a, Show a)
+
+foo :: PairDict a => a -> a -> String
+foo x y | x==y      = show x
+        | otherwise = show y
+
+-- In worker/wrapper we'd like to unbox the pair
+-- but not (Eq a) and (Show a)
+
+bar :: (a ~ b, Show a) => Int -> a -> (b, String)
+bar 0 x = (x, show x)
+bar n x = bar (n-1) x
diff --git a/testsuite/tests/stranal/should_compile/T23398.stderr b/testsuite/tests/stranal/should_compile/T23398.stderr
new file mode 100644
index 0000000000..84177a1424
--- /dev/null
+++ b/testsuite/tests/stranal/should_compile/T23398.stderr
@@ -0,0 +1,109 @@
+
+==================== Tidy Core ====================
+Result size of Tidy Core
+  = {terms: 76, types: 117, coercions: 4, joins: 0/0}
+
+-- RHS size: {terms: 18, types: 11, coercions: 0, joins: 0/0}
+T23398.$wfoo [InlPrag=[2]]
+  :: forall {a}. (Eq a, Show a) => a -> a -> String
+[GblId[StrictWorker([!, !])],
+ Arity=4,
+ Str=<SP(1C(1,C(1,L)),A)><SP(A,1C(1,L),A)><L><L>,
+ Unf=Unf{Src=<vanilla>, TopLvl=True,
+         Value=True, ConLike=True, WorkFree=True, Expandable=True,
+         Guidance=IF_ARGS [30 60 0 0] 120 0}]
+T23398.$wfoo
+  = \ (@a) (ww :: Eq a) (ww1 :: Show a) (eta :: a) (eta1 :: a) ->
+      case == @a ww eta eta1 of {
+        False -> show @a ww1 eta1;
+        True -> show @a ww1 eta
+      }
+
+-- RHS size: {terms: 12, types: 12, coercions: 0, joins: 0/0}
+foo [InlPrag=[2]] :: forall a. PairDict a => a -> a -> String
+[GblId,
+ Arity=3,
+ Str=<S!P(SP(SC(S,C(1,L)),A),SP(A,SC(S,L),A))><L><L>,
+ Unf=Unf{Src=StableSystem, TopLvl=True,
+         Value=True, ConLike=True, WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=3,unsat_ok=True,boring_ok=False)
+         Tmpl= \ (@a)
+                 ($d(%,%) [Occ=Once1!] :: PairDict a)
+                 (eta [Occ=Once1] :: a)
+                 (eta1 [Occ=Once1] :: a) ->
+                 case $d(%,%) of { (ww [Occ=Once1], ww1 [Occ=Once1]) ->
+                 T23398.$wfoo @a ww ww1 eta eta1
+                 }}]
+foo
+  = \ (@a) ($d(%,%) :: PairDict a) (eta :: a) (eta1 :: a) ->
+      case $d(%,%) of { (ww, ww1) -> T23398.$wfoo @a ww ww1 eta eta1 }
+
+Rec {
+-- RHS size: {terms: 21, types: 19, coercions: 3, joins: 0/0}
+T23398.$wbar [InlPrag=[2], Occ=LoopBreaker]
+  :: forall {a} {b}.
+     (a GHC.Prim.~# b, Show a) =>
+     GHC.Prim.Int# -> a -> (# b, String #)
+[GblId[StrictWorker([~, !])],
+ Arity=4,
+ Str=<L><SP(A,SC(S,L),A)><1L><L>,
+ Unf=OtherCon []]
+T23398.$wbar
+  = \ (@a)
+      (@b)
+      (ww :: a GHC.Prim.~# b)
+      ($dShow :: Show a)
+      (ww1 :: GHC.Prim.Int#)
+      (eta :: a) ->
+      case ww1 of ds {
+        __DEFAULT ->
+          T23398.$wbar
+            @a @b @~(ww :: a GHC.Prim.~# b) $dShow (GHC.Prim.-# ds 1#) eta;
+        0# -> (# eta `cast` (Sub ww :: a ~R# b), show @a $dShow eta #)
+      }
+end Rec }
+
+-- RHS size: {terms: 21, types: 32, coercions: 1, joins: 0/0}
+bar [InlPrag=[2]]
+  :: forall a b. (a ~ b, Show a) => Int -> a -> (b, String)
+[GblId,
+ Arity=4,
+ Str=<S!P(L)><SP(A,SC(S,L),A)><1!P(1L)><L>,
+ Cpr=1,
+ Unf=Unf{Src=StableSystem, TopLvl=True,
+         Value=True, ConLike=True, WorkFree=True, Expandable=True,
+         Guidance=ALWAYS_IF(arity=4,unsat_ok=True,boring_ok=False)
+         Tmpl= \ (@a)
+                 (@b)
+                 ($d~ [Occ=Once1!] :: a ~ b)
+                 ($dShow [Occ=Once1] :: Show a)
+                 (eta [Occ=Once1!] :: Int)
+                 (eta1 [Occ=Once1] :: a) ->
+                 case $d~ of { GHC.Types.Eq# ww ->
+                 case eta of { GHC.Types.I# ww1 [Occ=Once1] ->
+                 case T23398.$wbar @a @b @~(ww :: a GHC.Prim.~# b) $dShow ww1 eta1
+                 of
+                 { (# ww2 [Occ=Once1], ww3 [Occ=Once1] #) ->
+                 (ww2, ww3)
+                 }
+                 }
+                 }}]
+bar
+  = \ (@a)
+      (@b)
+      ($d~ :: a ~ b)
+      ($dShow :: Show a)
+      (eta :: Int)
+      (eta1 :: a) ->
+      case $d~ of { GHC.Types.Eq# ww ->
+      case eta of { GHC.Types.I# ww1 ->
+      case T23398.$wbar @a @b @~(ww :: a GHC.Prim.~# b) $dShow ww1 eta1
+      of
+      { (# ww2, ww3 #) ->
+      (ww2, ww3)
+      }
+      }
+      }
+
+
+
diff --git a/testsuite/tests/stranal/should_compile/all.T b/testsuite/tests/stranal/should_compile/all.T
index 4dbe61a300..145bc0eb9c 100644
--- a/testsuite/tests/stranal/should_compile/all.T
+++ b/testsuite/tests/stranal/should_compile/all.T
@@ -93,3 +93,4 @@ test('T22039', normal, compile, [''])
 test('T22388', [ grep_errmsg(r'^\S+\$w\S+') ], compile, ['-dsuppress-uniques -ddump-simpl'])
 # T22997: Just a panic that should not happen
 test('T22997', normal, compile, [''])
+test('T23398', normal, compile, ['-dsuppress-uniques -ddump-simpl -dno-typeable-binds'])