Make withDict opaque to the specialiser

As pointed out in #21575, it is not sufficient to set withDict to inline
after the typeclass specialiser, because we might inline withDict in one
module and then import it in another, and we run into the same problem.
This means we could still end up with incorrect runtime results because
the typeclass specialiser would assume that distinct typeclass evidence
terms at the same type are equal, when this is not necessarily the case
when using withDict.

Instead, this patch introduces a new magicId, 'nospec', which is only
inlined in CorePrep. We make use of it in the definition of withDict
to ensure that the typeclass specialiser does not common up distinct
typeclass evidence terms.

Fixes #21575

8 files changed, 193 insertions, 46 deletions

diff --git a/compiler/GHC/Builtin/Names.hs b/compiler/GHC/Builtin/Names.hs
index 115a7f53f4..f97103a90f 100644
--- a/compiler/GHC/Builtin/Names.hs
+++ b/compiler/GHC/Builtin/Names.hs
@@ -2332,6 +2332,9 @@ runRWKey                      = mkPreludeMiscIdUnique 107
 traceKey :: Unique
 traceKey                      = mkPreludeMiscIdUnique 108
 
+nospecIdKey :: Unique
+nospecIdKey                   = mkPreludeMiscIdUnique 109
+
 inlineIdKey, noinlineIdKey :: Unique
 inlineIdKey                   = mkPreludeMiscIdUnique 120
 -- see below
diff --git a/compiler/GHC/CoreToStg/Prep.hs b/compiler/GHC/CoreToStg/Prep.hs
index d03df0eedb..3078d26969 100644
--- a/compiler/GHC/CoreToStg/Prep.hs
+++ b/compiler/GHC/CoreToStg/Prep.hs
@@ -123,26 +123,24 @@ The goal of this pass is to prepare for code generation.
     We want curried definitions for all of these in case they
     aren't inlined by some caller.
 
-9.  Replace (lazy e) by e.  See Note [lazyId magic] in GHC.Types.Id.Make
-    Also replace (noinline e) by e.
+ 9. Convert bignum literals into their core representation.
 
-10. Convert bignum literals into their core representation.
-
-11. Uphold tick consistency while doing this: We move ticks out of
+10. Uphold tick consistency while doing this: We move ticks out of
     (non-type) applications where we can, and make sure that we
     annotate according to scoping rules when floating.
 
-12. Collect cost centres (including cost centres in unfoldings) if we're in
+11. Collect cost centres (including cost centres in unfoldings) if we're in
     profiling mode. We have to do this here beucase we won't have unfoldings
     after this pass (see `trimUnfolding` and Note [Drop unfoldings and rules].
 
-13. Eliminate case clutter in favour of unsafe coercions.
+12. Eliminate case clutter in favour of unsafe coercions.
     See Note [Unsafe coercions]
 
-14. Eliminate some magic Ids, specifically
+13. Eliminate some magic Ids, specifically
      runRW# (\s. e)  ==>  e[readWorldId/s]
-             lazy e  ==>  e
+             lazy e  ==>  e (see Note [lazyId magic] in GHC.Types.Id.Make)
          noinline e  ==>  e
+           nospec e  ==>  e
      ToDo:  keepAlive# ...
     This is done in cpeApp
 
@@ -1052,6 +1050,8 @@ cpeApp top_env expr
             -- See Note [lazyId magic] in GHC.Types.Id.Make
        || f `hasKey` noinlineIdKey      -- Replace (noinline a) with a
             -- See Note [noinlineId magic] in GHC.Types.Id.Make
+       || f `hasKey` nospecIdKey        -- Replace (nospec a) with a
+            -- See Note [nospecId magic] in GHC.Types.Id.Make
 
         -- Consider the code:
         --
diff --git a/compiler/GHC/Tc/Instance/Class.hs b/compiler/GHC/Tc/Instance/Class.hs
index f2ffeea8c6..2993d02ab6 100644
--- a/compiler/GHC/Tc/Instance/Class.hs
+++ b/compiler/GHC/Tc/Instance/Class.hs
@@ -33,6 +33,7 @@ import GHC.Types.SafeHaskell
 import GHC.Types.Name   ( Name, pprDefinedAt )
 import GHC.Types.Var.Env ( VarEnv )
 import GHC.Types.Id
+import GHC.Types.Id.Make ( nospecId )
 import GHC.Types.Var
 
 import GHC.Core.Predicate
@@ -43,9 +44,7 @@ import GHC.Core.DataCon
 import GHC.Core.TyCon
 import GHC.Core.Class
 
-import GHC.Core ( Expr(Var, App, Cast, Let), Bind (NonRec) )
-import GHC.Types.Basic
-import GHC.Types.SourceText
+import GHC.Core ( Expr(Var, App, Cast, Type) )
 
 import GHC.Utils.Outputable
 import GHC.Utils.Panic
@@ -455,20 +454,26 @@ matchWithDict [cls, mty]
   = do { sv <- mkSysLocalM (fsLit "withDict_s") Many mty
        ; k  <- mkSysLocalM (fsLit "withDict_k") Many (mkInvisFunTyMany cls openAlphaTy)
 
-       ; let evWithDict_type = mkSpecForAllTys [runtimeRep1TyVar, openAlphaTyVar] $
-                               mkVisFunTysMany [mty, mkInvisFunTyMany cls openAlphaTy] openAlphaTy
-
-       ; wd_id <- mkSysLocalM (fsLit "withDict_wd") Many evWithDict_type
-       ; let wd_id' = wd_id `setInlinePragma` inlineAfterSpecialiser
-
        -- Given co2 : mty ~N# inst_meth_ty, construct the method of
        -- the WithDict dictionary:
-       -- \@(r : RuntimeRep) @(a :: TYPE r) (sv : mty) (k :: cls => a) -> k (sv |> (sub co; sym co2))
+       --
+       --   \@(r :: RuntimeRep) @(a :: TYPE r) (sv :: mty) (k :: cls => a) ->
+       --     nospec @(cls => a) k (sv |> (sub co ; sym co2))
+       --
+       -- where  nospec :: forall a. a -> a  ensures that the typeclass specialiser
+       -- doesn't attempt to common up this evidence term with other evidence terms
+       -- of the same type.
+       --
+       -- See (WD6) in Note [withDict], and Note [nospecId magic] in GHC.Types.Id.Make.
        ; let evWithDict co2 =
-               let wd_rhs = mkCoreLams [ runtimeRep1TyVar, openAlphaTyVar, sv, k ] $
-                            Var k `App` Cast (Var sv) (mkTcTransCo (mkTcSubCo co2) (mkTcSymCo co))
-               in Let (NonRec wd_id' wd_rhs) (Var wd_id')
-         -- Why a Let?  See (WD6) in Note [withDict]
+               mkCoreLams [ runtimeRep1TyVar, openAlphaTyVar, sv, k ] $
+                 Var nospecId
+                   `App`
+                 (Type $ mkInvisFunTyMany cls openAlphaTy)
+                   `App`
+                 Var k
+                   `App`
+                 (Var sv `Cast` mkTcTransCo (mkTcSubCo co2) (mkTcSymCo co))
 
        ; tc <- tcLookupTyCon withDictClassName
        ; let Just withdict_data_con
@@ -486,12 +491,6 @@ matchWithDict [cls, mty]
 matchWithDict _
   = return NoInstance
 
-inlineAfterSpecialiser :: InlinePragma
--- Do not inline before the specialiser; but do so afterwards
--- See (WD6) in Note [withDict]
-inlineAfterSpecialiser = alwaysInlinePragma `setInlinePragmaActivation`
-                         ActiveAfter NoSourceText 2
-
 {-
 Note [withDict]
 ~~~~~~~~~~~~~~~
@@ -536,7 +535,7 @@ as if the following instance declaration existed:
 
 instance (mty ~# inst_meth_ty) => WithDict (C t1..tn) mty where
   withDict = \@{rr} @(r :: TYPE rr) (sv :: mty) (k :: C t1..tn => r) ->
-    k (sv |> (sub co2; sym co))
+    k (sv |> (sub co2 ; sym co))
 
 That is, it matches on the first (constraint) argument of C; if C is
 a single-method class, the instance "fires" and emits an equality
@@ -582,28 +581,55 @@ Some further observations about `withDict`:
 
 (WD5) In earlier implementations, `withDict` was implemented as an identifier
       with special handling during either constant-folding or desugaring.
-      The current approach is more robust, previously the type of `withDict`
+      The current approach is more robust: previously, the type of `withDict`
       did not have a type-class constraint and was overly polymorphic.
       See #19915.
 
-(WD6) In fact we desugar `withDict @(C t_1 ... t_n) @mty @{rr} @r` to
-         let wd = \sv k -> k (sv |> co)
-             {-# INLINE [2] #-}
-         in wd
-      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
+(WD6) In fact, we desugar `withDict @cls @mty @{rr} @r` to
+
+         \@(r :: RuntimeRep) @(a :: TYPE r) (sv :: mty) (k :: cls => a) ->
+           nospec @(cls => a) k (sv |> (sub co2 ; sym co)))
+
+      That is, we cast the method using a coercion, and apply k to it.
+      However, we use the 'nospec' magicId (see Note [nospecId magic] in GHC.Types.Id.Make)
+      to ensure that the typeclass specialiser doesn't incorrectly common-up distinct
+      evidence terms. 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
+
+      where k1, k2 :: C T -> blah.  If we desugared withDict naively, we'd 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,
+
+      and the Specialiser would assume that those arguments (of type `C T`) are
+      the same. It would then specialise `k` for that type, and then call the same,
       specialised function from both call sites.  #21575 is a concrete case in point.
 
-      Solution: delay inlining `withDict` until after the specialiser; that is,
-      until Phase 2.  This is not a Final Solution -- seee #21575 "Alas..".
+      To avoid this, we need to stop the typeclass specialiser from seeing this
+      structure, by using nospec. This function is inlined only in CorePrep; crucially
+      this means that it still appears in interface files, so that the desugaring of
+      withDict remains opaque to the typeclass specialiser across modules.
+      This means the specialiser will always see instead:
+
+          nospec @(cls => a) k (A |> co1)
+          nospec @(cls => a) k (B |> co2)
+
+      Why does this work? Recall that nospec is not an overloaded function;
+      it has the type
+
+        nospec :: forall a. a -> a
+
+      This means that there is nothing for the specialiser to do with function calls
+      such as
+
+        nospec @(cls => a) k (A |> co)
+
+      as the specialiser only looks at calls of the form `f dict` for an
+      overloaded function `f` (e.g. with a type such as `f :: Eq a => ...`).
+
+      See test-case T21575b.
 
 -}
 
diff --git a/compiler/GHC/Types/Id/Make.hs b/compiler/GHC/Types/Id/Make.hs
index 9628bea733..e31460de7c 100644
--- a/compiler/GHC/Types/Id/Make.hs
+++ b/compiler/GHC/Types/Id/Make.hs
@@ -32,7 +32,7 @@ module GHC.Types.Id.Make (
         voidPrimId, voidArgId,
         nullAddrId, seqId, lazyId, lazyIdKey,
         coercionTokenId, coerceId,
-        proxyHashId, noinlineId, noinlineIdName,
+        proxyHashId, noinlineId, noinlineIdName, nospecId, nospecIdName,
         coerceName, leftSectionName, rightSectionName,
     ) where
 
@@ -159,7 +159,7 @@ wiredInIds
   ++ errorIds           -- Defined in GHC.Core.Make
 
 magicIds :: [Id]    -- See Note [magicIds]
-magicIds = [lazyId, oneShotId, noinlineId]
+magicIds = [lazyId, oneShotId, noinlineId, nospecId]
 
 ghcPrimIds :: [Id]  -- See Note [ghcPrimIds (aka pseudoops)]
 ghcPrimIds
@@ -1401,10 +1401,11 @@ leftSectionName   = mkWiredInIdName gHC_PRIM  (fsLit "leftSection")    leftSecti
 rightSectionName  = mkWiredInIdName gHC_PRIM  (fsLit "rightSection")   rightSectionKey    rightSectionId
 
 -- Names listed in magicIds; see Note [magicIds]
-lazyIdName, oneShotName, noinlineIdName :: Name
+lazyIdName, oneShotName, noinlineIdName, nospecIdName :: Name
 lazyIdName        = mkWiredInIdName gHC_MAGIC (fsLit "lazy")           lazyIdKey          lazyId
 oneShotName       = mkWiredInIdName gHC_MAGIC (fsLit "oneShot")        oneShotKey         oneShotId
 noinlineIdName    = mkWiredInIdName gHC_MAGIC (fsLit "noinline")       noinlineIdKey      noinlineId
+nospecIdName      = mkWiredInIdName gHC_MAGIC (fsLit "nospec")         nospecIdKey        nospecId
 
 ------------------------------------------------
 proxyHashId :: Id
@@ -1472,6 +1473,12 @@ noinlineId = pcMiscPrelId noinlineIdName ty info
     info = noCafIdInfo
     ty  = mkSpecForAllTys [alphaTyVar] (mkVisFunTyMany alphaTy alphaTy)
 
+nospecId :: Id -- See Note [nospecId magic]
+nospecId = pcMiscPrelId nospecIdName ty info
+  where
+    info = noCafIdInfo
+    ty  = mkSpecForAllTys [alphaTyVar] (mkVisFunTyMany alphaTy alphaTy)
+
 oneShotId :: Id -- See Note [The oneShot function]
 oneShotId = pcMiscPrelId oneShotName ty info
   where
@@ -1727,6 +1734,19 @@ special case to the demand analyser to address #16588. However, the special
 case seemed like a large and expensive hammer to address a rare case and
 consequently we rather opted to use a more minimal solution.
 
+Note [nospecId magic]
+~~~~~~~~~~~~~~~~~~~~~
+The 'nospec' magic Id is used to ensure to make a value opaque to the typeclass
+specialiser. In CorePrep, we inline 'nospec', turning (nospec e) into e.
+Note that this happens *after* unfoldings are exposed in the interface file.
+This is crucial: otherwise, we could import an unfolding in which
+'nospec' has been inlined (= erased), and we would lose the benefit.
+
+'nospec' is used in the implementation of 'withDict': we insert 'nospec'
+so that the typeclass specialiser doesn't assume any two evidence terms
+of the same type are equal. See Note [withDict] in GHC.Tc.Instance.Class,
+and see test case T21575b for an example.
+
 Note [The oneShot function]
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~
 In the context of making left-folds fuse somewhat okish (see ticket #7994
diff --git a/testsuite/tests/simplCore/should_run/T21575b.hs b/testsuite/tests/simplCore/should_run/T21575b.hs
new file mode 100644
index 0000000000..da6362edaa
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T21575b.hs
@@ -0,0 +1,69 @@
+{-# LANGUAGE UndecidableInstances #-}
+
+module Main (main) where
+
+import Control.Monad (unless)
+import qualified Data.Map as M
+import Data.Map (Map)
+import T21575b_aux (Given(..), give, unsafeGive)
+
+main :: IO ()
+main = do
+  testCase "Normal"
+    (give Normal (toJSON (Foo Bar)))
+    (unsafeGive Normal (toJSON (Foo Bar)))
+    (Object (M.fromList [("Foo",String "Bar")]))
+  testCase "ViaShow"
+    (give ViaShow (toJSON (Foo Bar)))
+    (unsafeGive ViaShow (toJSON (Foo Bar)))
+    (Object (M.fromList [("Foo",String "SHOWBAR")]))
+
+-----
+
+testCase :: (Eq a, Show a) => String -> a -> a -> a -> IO ()
+testCase str with_give with_unsafeGive expected =
+  putStrLn $ unlines
+      [ str ++ ":"
+      , "    withDict: " ++ show with_give
+      , "unsafeCoerce: " ++ show with_unsafeGive
+      , "    expected: " ++ show expected
+      , ""
+      ]
+
+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/T21575b.stdout b/testsuite/tests/simplCore/should_run/T21575b.stdout
new file mode 100644
index 0000000000..2cfdf0dcac
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T21575b.stdout
@@ -0,0 +1,12 @@
+Normal:
+    withDict: Object (fromList [("Foo",String "Bar")])
+unsafeCoerce: Object (fromList [("Foo",String "Bar")])
+    expected: Object (fromList [("Foo",String "Bar")])
+
+
+ViaShow:
+    withDict: Object (fromList [("Foo",String "SHOWBAR")])
+unsafeCoerce: Object (fromList [("Foo",String "SHOWBAR")])
+    expected: Object (fromList [("Foo",String "SHOWBAR")])
+
+
diff --git a/testsuite/tests/simplCore/should_run/T21575b_aux.hs b/testsuite/tests/simplCore/should_run/T21575b_aux.hs
new file mode 100644
index 0000000000..556ef26307
--- /dev/null
+++ b/testsuite/tests/simplCore/should_run/T21575b_aux.hs
@@ -0,0 +1,16 @@
+module T21575b_aux ( Given(..), give, unsafeGive ) where
+
+import GHC.Exts
+  ( withDict )
+import Unsafe.Coerce
+  ( unsafeCoerce )
+
+class Given a where
+  given :: a
+
+give, unsafeGive :: forall a r. a -> (Given a => r) -> r
+give = withDict @(Given a) @a
+
+unsafeGive a k = unsafeCoerce (Gift k :: Gift a r) a
+
+newtype Gift a r = Gift (Given a => r)
diff --git a/testsuite/tests/simplCore/should_run/all.T b/testsuite/tests/simplCore/should_run/all.T
index bebd839724..a887d5cedb 100644
--- a/testsuite/tests/simplCore/should_run/all.T
+++ b/testsuite/tests/simplCore/should_run/all.T
@@ -105,4 +105,5 @@ 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'])
+test('T21575b', [], multimod_compile_and_run, ['T21575b', '-O'])
 test('T20836', normal, compile_and_run, ['-O0']) # Should not time out; See #20836