Fix bad interaction between withDict and the Specialiserwip/T21575

This MR fixes a bad bug, where the withDict was inlined too
vigorously, which in turn made the type-class Specialiser generate
a bogus specialisation, because it saw the same overloaded function
applied to two /different/ dictionaries.

Solution: inline `withDict` later.  See (WD8) of Note [withDict]
in GHC.HsToCore.Expr

See #21575, which is fixed by this change.

7 files changed, 201 insertions, 51 deletions

diff --git a/compiler/GHC/Core/Coercion.hs b/compiler/GHC/Core/Coercion.hs
index 235e8c65fb..98aa337026 100644
--- a/compiler/GHC/Core/Coercion.hs
+++ b/compiler/GHC/Core/Coercion.hs
@@ -1708,9 +1708,8 @@ has to do in its full generality.  See #18413.
 %************************************************************************
 -}
 
--- | If @co :: T ts ~ rep_ty@ then:
---
--- > instNewTyCon_maybe T ts = Just (rep_ty, co)
+-- | If `instNewTyCon_maybe T ts = Just (rep_ty, co)`
+--   then `co :: T ts ~R# rep_ty`
 --
 -- Checks for a newtype, and for being saturated
 instNewTyCon_maybe :: TyCon -> [Type] -> Maybe (Type, Coercion)
diff --git a/compiler/GHC/Core/TyCon.hs b/compiler/GHC/Core/TyCon.hs
index e50751c2a0..6514b68aaf 100644
--- a/compiler/GHC/Core/TyCon.hs
+++ b/compiler/GHC/Core/TyCon.hs
@@ -1184,7 +1184,9 @@ data AlgTyConRhs
                         -- See Note [Newtype eta]
         nt_co :: CoAxiom Unbranched,
                              -- The axiom coercion that creates the @newtype@
-                             -- from the representation 'Type'.
+                             -- from the representation 'Type'.  The axiom witnesses
+                             -- a representational coercion:
+                             --   nt_co :: N ty1 ~R# rep_tys
 
                              -- See Note [Newtype coercions]
                              -- Invariant: arity = #tvs in nt_etad_rhs;
diff --git a/compiler/GHC/HsToCore/Expr.hs b/compiler/GHC/HsToCore/Expr.hs
index 88dc4f7d45..6b8ced95ad 100644
--- a/compiler/GHC/HsToCore/Expr.hs
+++ b/compiler/GHC/HsToCore/Expr.hs
@@ -1140,15 +1140,26 @@ ds_withDict wrapped_ty
     -- `meth_tvs = a_1 ... a_n` and `co` is a newtype coercion between
     -- `C` and `meth_ty`.
   , Just (inst_meth_ty, co) <- instNewTyCon_maybe dict_tc dict_args
+    --    co :: C t1 ..tn ~R# st
     -- Check that `st` is equal to `meth_ty[t_i/a_i]`.
   , st `eqType` inst_meth_ty
-  = do { sv <- newSysLocalDs mult1 st
+ = do { sv <- newSysLocalDs mult1 st
        ; k  <- newSysLocalDs mult2 dt_to_r
-       ; pure $ mkLams [sv, k] $ Var k `App` Cast (Var sv) (mkSymCo co) }
+       ; let wd_rhs = mkLams [sv, k] $ Var k `App` Cast (Var sv) (mkSymCo co)
+       ; wd_id <- newSysLocalDs Many (exprType wd_rhs)
+       ; let wd_id' = wd_id `setInlinePragma` inlineAfterSpecialiser
+       ; pure $ Let (NonRec wd_id' wd_rhs) (Var wd_id') }
+         -- Why a Let?  See (WD8) in Note [withDict]
 
   | otherwise
   = errDsCoreExpr (DsInvalidInstantiationDictAtType wrapped_ty)
 
+inlineAfterSpecialiser :: InlinePragma
+-- Do not inline before the specialiser; but do so afterwards
+-- See (WD8) in Note [withDict]
+inlineAfterSpecialiser = alwaysInlinePragma `setInlinePragmaActivation`
+                         ActiveAfter NoSourceText 2
+
 {- Note [withDict]
 ~~~~~~~~~~~~~~~~~~
 The identifier `withDict` is just a place-holder, which is used to
@@ -1202,7 +1213,7 @@ Where:
   That is, C must be a class with exactly one method and no superclasses.
 
 * The `mtype` argument to withDict must be equal to `meth_type[t_i/a_i]`,
-  which is instantied type of C's method.
+  which is instantiated type of C's method.
 
 * `co` is a newtype coercion that, when applied to `t_1 ... t_n`, coerces from
   `C t_1 ... t_n` to `mtype`. This coercion is guaranteed to exist by virtue of
@@ -1213,66 +1224,86 @@ These requirements are implemented in the guards in ds_withDict's definition.
 
 Some further observations about `withDict`:
 
-* Every use of `withDict` must be instantiated at a /particular/ class C.
-  It's a bit like representation polymorphism: we don't allow class-polymorphic
-  calls of `withDict`. We check this in the desugarer -- and then we
-  can immediately replace this invocation of `withDict` with appropriate
-  class-specific Core code.
+(WD1) Every use of `withDict` must be instantiated at a /particular/ class C.
+      It's a bit like representation polymorphism: we don't allow class-polymorphic
+      calls of `withDict`. We check this in the desugarer -- and then we
+      can immediately replace this invocation of `withDict` with appropriate
+      class-specific Core code.
+
+(WD2) The `dt` in the type of withDict must be explicitly instantiated with
+      visible type application, as invoking `withDict` would be ambiguous
+      otherwise.
+
+      For examples of how `withDict` is used in the `base` library, see `withSNat`
+      in GHC.TypeNats, as well as `withSChar` and `withSSymbol` in GHC.TypeLits.
+
+(WD3) The `r` is representation-polymorphic, to support things like
+      `withTypeable` in `Data.Typeable.Internal`.
 
-* The `dt` in the type of withDict must be explicitly instantiated with
-  visible type application, as invoking `withDict` would be ambiguous
-  otherwise.
+(WD4) As an alternative to `withDict`, one could define functions like `withT`
+      above in terms of `unsafeCoerce`. This is more error-prone, however.
 
-* For examples of how `withDict` is used in the `base` library, see `withSNat`
-  in GHC.TypeNats, as well as `withSChar` and `withSSymbol` in GHC.TypeLits.
+(WD5) In order to define things like `reifySymbol` below:
 
-* The `r` is representation-polymorphic,
-  to support things like `withTypeable` in `Data.Typeable.Internal`.
+        reifySymbol :: forall r. String -> (forall (n :: Symbol). KnownSymbol n => r) -> r
 
-* As an alternative to `withDict`, one could define functions like `withT`
-  above in terms of `unsafeCoerce`. This is more error-prone, however.
+      `withDict` needs to be instantiated with `Any`, like so:
 
-* In order to define things like `reifySymbol` below:
+        reifySymbol n k = withDict @String @(KnownSymbol Any) @r n (k @Any)
 
-    reifySymbol :: forall r. String -> (forall (n :: Symbol). KnownSymbol n => r) -> r
+      The use of `Any` is explained in Note [NOINLINE someNatVal] in
+      base:GHC.TypeNats.
 
-  `withDict` needs to be instantiated with `Any`, like so:
+(WD6) The only valid way to apply `withDict` is as described above. Applying
+      `withDict` in any other way will result in a non-recoverable error during
+      desugaring. In other words, GHC will never execute the `withDict` function
+      in compiled code.
 
-    reifySymbol n k = withDict @String @(KnownSymbol Any) @r n (k @Any)
+      In theory, this means that we don't need to define a binding for `withDict`
+      in GHC.Magic.Dict. In practice, we define a binding anyway, for two reasons:
 
-  The use of `Any` is explained in Note [NOINLINE someNatVal] in
-  base:GHC.TypeNats.
+        - To give it Haddocks, and
+        - To define the type of `withDict`, which GHC can find in
+          GHC.Magic.Dict.hi.
 
-* The only valid way to apply `withDict` is as described above. Applying
-  `withDict` in any other way will result in a non-recoverable error during
-  desugaring. In other words, GHC will never execute the `withDict` function
-  in compiled code.
+      Because we define a binding for `withDict`, we have to provide a right-hand
+      side for its definition. We somewhat arbitrarily choose:
 
-  In theory, this means that we don't need to define a binding for `withDict`
-  in GHC.Magic.Dict. In practice, we define a binding anyway, for two reasons:
+        withDict = panicError "Non rewritten withDict"#
 
-    - To give it Haddocks, and
-    - To define the type of `withDict`, which GHC can find in
-      GHC.Magic.Dict.hi.
+      This should never be reachable anyway, but just in case ds_withDict fails
+      to rewrite away `withDict`, this ensures that the program won't get very far.
 
-  Because we define a binding for `withDict`, we have to provide a right-hand
-  side for its definition. We somewhat arbitrarily choose:
+(WD7) One could conceivably implement this special case for `withDict` as a
+      constant-folding rule instead of during desugaring. We choose not to do so
+      for the following reasons:
 
-    withDict = panicError "Non rewritten withDict"#
+      - Having a constant-folding rule would require that `withDict`'s definition
+        be wired in to the compiler so as to prevent `withDict` from inlining too
+        early. Implementing the special case in the desugarer, on the other hand,
+        only requires that `withDict` be known-key.
 
-  This should never be reachable anyway, but just in case ds_withDict fails
-  to rewrite away `withDict`, this ensures that the program won't get very far.
+      - If the constant-folding rule were to fail, we want to throw a compile-time
+        error, which is trickier to do with the way that GHC.Core.Opt.ConstantFold
+        is set up.
 
-* One could conceivably implement this special case for `withDict` as a
-  constant-folding rule instead of during desugaring. We choose not to do so
-  for the following reasons:
+(WD8) In fact we desugar `withDict @{rr} @mtype @(C t_1 ... t_n) @r` to
+         let wd = \sv k -> k (sv |> co)
+             {-# INLINE [2] #-}
+         in wd
 
-  - Having a constant-folding rule would require that `withDict`'s definition
-    be wired in to the compiler so as to prevent `withDict` from inlining too
-    early. Implementing the special case in the desugarer, on the other hand,
-    only requires that `withDict` be known-key.
+      The local `let` and INLINE pragma delays inlining `wd` until after the
+      type-class Specialiser has run.  This is super important. Suppose we
+      have calls
+          withDict A k
+          withDict B k
+      where k1, k2 :: C T -> blah.  If we inline those withDict calls we'll get
+          k (A |> co1)
+          k (B |> co2)
+      and the Specialiser will assume that those arguments (of type `C T`) are
+      the same, will specialise `k` for that type, and will call the same,
+      specialised function from both call sites.  #21575 is a concrete case in point.
 
-  - If the constant-folding rule were to fail, we want to throw a compile-time
-    error, which is trickier to do with the way that GHC.Core.Opt.ConstantFold
-    is set up.
+      Solution: delay inlining `withDict` until after the specialiser; that is,
+      until Phase 2.  This is not a Final Solution -- seee #21575 "Alas..".
 -}
diff --git a/libraries/base/Unsafe/Coerce.hs b/libraries/base/Unsafe/Coerce.hs
index ec5b793f87..fd45217fef 100644
--- a/libraries/base/Unsafe/Coerce.hs
+++ b/libraries/base/Unsafe/Coerce.hs
@@ -175,6 +175,12 @@ several ways
      no longer does that.  So we need a separate map/unsafeCoerce
      RULE, in this module.
 
+     Adding these RULES means we must delay inlinine unsafeCoerce
+     until the RULES have had a chance to fire; hence the INLINE[1]
+     pragma on unsafeCoerce.  (Side note: this has the coincidental
+     benefit of making the unsafeCoerce-based version of the `reflection`
+     library work -- see #21575.)
+
 There are yet more wrinkles
 
 (U9) unsafeCoerce works only over types of kind `Type`.
@@ -236,6 +242,9 @@ unsafeEqualityProof = case unsafeEqualityProof @a @b of UnsafeRefl -> UnsafeRefl
 --   case unsafeEqualityProof of UnsafeRefl -> case blah of ...
 --
 -- which is definitely better.
+--
+-- Why delay inlining to Phase 1?  Because of the RULES for map/unsafeCoerce;
+-- see (U8) in Note [Implementing unsafeCoerce]
 
 -- | Coerce a value from one type to another, bypassing the type-checker.
 --
diff --git a/testsuite/tests/simplCore/should_run/T21575.hs b/testsuite/tests/simplCore/should_run/T21575.hs
new file mode 100644
index 0000000000..976483f963
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T21575.hs
@@ -0,0 +1,107 @@
+{-# LANGUAGE CPP #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+-- 0 => use unsafeCoerce
+-- 1 => use withDict
+#define WITH_DICT 1
+
+module Main (main) where
+
+import Control.Monad (unless)
+import qualified Data.Map as M
+import Data.Map (Map)
+
+#if WITH_DICT
+import GHC.Exts (withDict)
+#else
+import Unsafe.Coerce (unsafeCoerce)
+#endif
+
+main :: IO ()
+main = do
+  testCase (give Normal (toJSON (Foo Bar)))
+           (Object (M.fromList [("Foo",String "Bar")]))
+  testCase (give ViaShow (toJSON (Foo Bar)))
+           (Object (M.fromList [("Foo",String "SHOWBAR")]))
+  putStrLn "All tests passed!"
+
+{-
+toJSONBar :: Given Style => Bar -> Value
+
+  give Normal (\gd -> toJSONBar gd e)
+  --> withDict @Style @(Given Style) Normal (toJSON e)
+  --> toJSONBar ((Normal |> co) :: Given Style) e
+
+  give Normal (\gd -> toJSONBar gd e')
+  --> toJSONBar ((ViaShow |> co) :: Given Style) e'
+
+--------- With new cast ------------
+
+  give Normal (\gd -> toJSONBar gd e)
+  --> withDict @Style @(Given Style) Normal (\gd -> toJSONBar gd e)
+  --> ((\gd -> toJSONBar gd e) |> co) Normal
+  --> (\gd' -> toJSonBar (gd' |> sym (co[1])) e) Normal
+  --> toJSONBar (Normal |> co') e   -- Boo!
+
+-}
+
+testCase :: (Eq a, Show a) => a -> a -> IO ()
+testCase expected actual =
+  unless (expected == actual) $
+    error $ unlines
+      [ ""
+      , "Expected: " ++ show expected
+      , "Actual:   " ++ show actual
+      ]
+
+class Given a where
+  given :: a
+
+give :: forall a r. a -> (Given a => r) -> r
+#if WITH_DICT
+give = withDict @a @(Given a)
+#else
+give a k = unsafeCoerce (Gift k :: Gift a r) a
+
+newtype Gift a r = Gift (Given a => r)
+#endif
+
+data Foo = Foo Bar
+
+instance Show Foo where
+  show _ = "SHOWFOO"
+
+data Bar = Bar | BarBar
+
+instance Show Bar where
+  show _ = "SHOWBAR"
+
+----------------------------------------------------------------------------
+-- ToJSON instances
+----------------------------------------------------------------------------
+
+instance Given Style => ToJSON Foo where
+  toJSON (Foo x) = Object $ M.singleton "Foo" (toJSON x)
+
+instance Given Style => ToJSON Bar where
+  toJSON x = case given of
+    Normal -> String $ case x of
+                Bar    -> "Bar"
+                BarBar -> "BarBar"
+    ViaShow -> String $ show x
+
+data Style = Normal | ViaShow
+
+----------------------------------------------------------------------------
+-- Minimized aeson
+----------------------------------------------------------------------------
+
+class ToJSON a where
+  toJSON :: a -> Value
+
+data Value
+  = Object !(Map String Value)
+  | String !String
+  deriving (Eq, Show)
diff --git a/testsuite/tests/simplCore/should_run/T21575.stdout b/testsuite/tests/simplCore/should_run/T21575.stdout
new file mode 100644
index 0000000000..14332104f6
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T21575.stdout
@@ -0,0 +1 @@
+All tests passed!
diff --git a/testsuite/tests/simplCore/should_run/all.T b/testsuite/tests/simplCore/should_run/all.T
index 52a1af40a8..53bcde5169 100644
--- a/testsuite/tests/simplCore/should_run/all.T
+++ b/testsuite/tests/simplCore/should_run/all.T
@@ -102,3 +102,4 @@ test('T20203', normal, compile, ['-O -dsuppress-all -dsuppress-uniques -dno-type
 test('T19313', normal, compile_and_run, [''])
 test('UnliftedArgRule', normal, compile_and_run, [''])
 test('T21229', normal, compile_and_run, ['-O'])
+test('T21575', normal, compile_and_run, ['-O'])