Refactor handling of SPECIALISE pragmas (Trac #5821)

The provoking cause was Trac #5821, which concerned
type families, but in fixing it I did the usual round
of tidying up and docmenting.

The main comment is now
    Note [Handling SPECIALISE pragmas]
in TcBinds.

It is "wrinkle 2" that fixes #5821.

compiler/deSugar/DsBinds.hs

@@ -389,42 +389,6 @@ gotten from the binding for fromT_1.
 
 It might be better to have just one level of AbsBinds, but that requires more
 thought!
-
-Note [Implementing SPECIALISE pragmas]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-Example:
-        f :: (Eq a, Ix b) => a -> b -> Bool
-        {-# SPECIALISE f :: (Ix p, Ix q) => Int -> (p,q) -> Bool #-}
-        f = <poly_rhs>
-
-From this the typechecker generates
-
-    AbsBinds [ab] [d1,d2] [([ab], f, f_mono, prags)] binds
-
-    SpecPrag (wrap_fn :: forall a b. (Eq a, Ix b) => XXX
-                      -> forall p q. (Ix p, Ix q) => XXX[ Int/a, (p,q)/b ])
-
-Note that wrap_fn can transform *any* function with the right type prefix
-    forall ab. (Eq a, Ix b) => XXX
-regardless of XXX.  It's sort of polymorphic in XXX.  This is
-useful: we use the same wrapper to transform each of the class ops, as
-well as the dict.
-
-From these we generate:
-
-    Rule:       forall p, q, (dp:Ix p), (dq:Ix q).
-                    f Int (p,q) dInt ($dfInPair dp dq) = f_spec p q dp dq
-
-    Spec bind:  f_spec = wrap_fn <poly_rhs>
-
-Note that
-
-  * The LHS of the rule may mention dictionary *expressions* (eg
-    $dfIxPair dp dq), and that is essential because the dp, dq are
-    needed on the RHS.
-
-  * The RHS of f_spec, <poly_rhs> has a *copy* of 'binds', so that it
-    can fully specialise it.
 -}
 
 ------------------------
@@ -432,7 +396,7 @@ dsSpecs :: CoreExpr     -- Its rhs
         -> TcSpecPrags
         -> DsM ( OrdList (Id,CoreExpr)  -- Binding for specialised Ids
                , [CoreRule] )           -- Rules for the Global Ids
--- See Note [Implementing SPECIALISE pragmas]
+-- See Note [Handling SPECIALISE pragmas] in TcBinds
 dsSpecs _ IsDefaultMethod = return (nilOL, [])
 dsSpecs poly_rhs (SpecPrags sps)
   = do { pairs <- mapMaybeM (dsSpec (Just poly_rhs)) sps
@@ -698,7 +662,7 @@ drop_dicts drops dictionary bindings on the LHS where possible.
          Here we want to end up with
             RULE forall d:Eq a.  f ($dfEqList d) = f_spec d
          Of course, the ($dfEqlist d) in the pattern makes it less likely
-         to match, but ther is no other way to get d:Eq a
+         to match, but there is no other way to get d:Eq a
 
    NB 2: We do drop_dicts *before* simplOptEpxr, so that we expect all
          the evidence bindings to be wrapped around the outside of the
@@ -714,7 +678,7 @@ drop_dicts drops dictionary bindings on the LHS where possible.
              useAbstractMonad :: MonadAbstractIOST m => m Int
          Here, deriving (MonadAbstractIOST (ReaderST s)) is a lot of code
          but the RHS uses no dictionaries, so we want to end up with
-             RULE forall s (d :: MonadBstractIOST (ReaderT s)).
+             RULE forall s (d :: MonadAbstractIOST (ReaderT s)).
                 useAbstractMonad (ReaderT s) d = $suseAbstractMonad s
 
    Trac #8848 is a good example of where there are some intersting

compiler/typecheck/TcBinds.hs

@@ -9,8 +9,9 @@
 
 module TcBinds ( tcLocalBinds, tcTopBinds, tcRecSelBinds,
                  tcHsBootSigs, tcPolyCheck,
-                 PragFun, tcSpecPrags, tcVectDecls, mkPragFun,
-                 TcSigInfo(..), TcSigFun,
+                 PragFun, tcSpecPrags, tcSpecWrapper,
+                 tcVectDecls, 
+                 TcSigInfo(..), TcSigFun, mkPragFun,
                  instTcTySig, instTcTySigFromId, findScopedTyVars,
                  badBootDeclErr, mkExport ) where
 
@@ -29,6 +30,7 @@ import TcHsType
 import TcPat
 import TcMType
 import ConLike
+import Inst( deeplyInstantiate )
 import FamInstEnv( normaliseType )
 import FamInst( tcGetFamInstEnvs )
 import Type( pprSigmaTypeExtraCts )
@@ -837,6 +839,135 @@ It also cleverly does an ambiguity check; for example, rejecting
 where F is a non-injective type function.
 -}
 
+--------------
+-- If typechecking the binds fails, then return with each
+-- signature-less binder given type (forall a.a), to minimise
+-- subsequent error messages
+recoveryCode :: [Name] -> TcSigFun -> TcM (LHsBinds TcId, [Id], TopLevelFlag)
+recoveryCode binder_names sig_fn
+  = do  { traceTc "tcBindsWithSigs: error recovery" (ppr binder_names)
+        ; poly_ids <- mapM mk_dummy binder_names
+        ; return (emptyBag, poly_ids, if all is_closed poly_ids
+                                      then TopLevel else NotTopLevel) }
+  where
+    mk_dummy name
+        | isJust (sig_fn name) = tcLookupId name        -- Had signature; look it up
+        | otherwise            = return (mkLocalId name forall_a_a)    -- No signature
+
+    is_closed poly_id = isEmptyVarSet (tyVarsOfType (idType poly_id))
+
+forall_a_a :: TcType
+forall_a_a = mkForAllTy openAlphaTyVar (mkTyVarTy openAlphaTyVar)
+
+
+{- *********************************************************************
+*                                                                      *
+                   Pragmas, including SPECIALISE
+*                                                                      *
+************************************************************************
+
+Note [Handling SPECIALISE pragmas]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+The basic idea is this:
+
+   f:: Num a => a -> b -> a
+   {-# SPECIALISE foo :: Int -> b -> Int #-}
+
+We check that 
+   (forall a. Num a => a -> a) 
+      is more polymorphic than 
+   Int -> Int
+(for which we could use tcSubType, but see below), generating a HsWrapper
+to connect the two, something like
+      wrap = /\b. <hole> Int b dNumInt
+This wrapper is put in the TcSpecPrag, in the ABExport record of
+the AbsBinds.
+
+
+        f :: (Eq a, Ix b) => a -> b -> Bool
+        {-# SPECIALISE f :: (Ix p, Ix q) => Int -> (p,q) -> Bool #-}
+        f = <poly_rhs>
+
+From this the typechecker generates
+
+    AbsBinds [ab] [d1,d2] [([ab], f, f_mono, prags)] binds
+
+    SpecPrag (wrap_fn :: forall a b. (Eq a, Ix b) => XXX
+                      -> forall p q. (Ix p, Ix q) => XXX[ Int/a, (p,q)/b ])
+
+From these we generate:
+
+    Rule:       forall p, q, (dp:Ix p), (dq:Ix q).
+                    f Int (p,q) dInt ($dfInPair dp dq) = f_spec p q dp dq
+
+    Spec bind:  f_spec = wrap_fn <poly_rhs>
+
+Note that
+
+  * The LHS of the rule may mention dictionary *expressions* (eg
+    $dfIxPair dp dq), and that is essential because the dp, dq are
+    needed on the RHS.
+
+  * The RHS of f_spec, <poly_rhs> has a *copy* of 'binds', so that it
+    can fully specialise it.
+
+
+
+From the TcSpecPrag, in DsBinds we generate a binding for f_spec and a RULE:
+
+   f_spec :: Int -> b -> Int
+   f_spec = wrap<f rhs>
+
+   RULE: forall b (d:Num b). f b d = f_spec b
+
+The RULE is generated by taking apart the HsWrapper, which is a little
+delicate, but works.  
+
+Some wrinkles
+
+1. We don't use full-on tcSubType, because that does co and contra
+   variance and that in turn will generate too complex a LHS for the
+   RULE.  So we use a single invocation of deeplySkolemise /
+   deeplyInstantiate in tcSpecWrapper.  (Actually I think that even
+   the "deeply" stuff may be too much, because it introduces lambdas,
+   though I think it can be made to work without too much trouble.)
+
+2. We need to take care with type families (Trac #5821).  Consider
+      type instance F Int = Bool
+      f :: Num a => a -> F a
+      {-# SPECIALISE foo :: Int -> Bool #-}
+
+  We *could* try to generate an f_spec with precisely the declared type:
+      f_spec :: Int -> Bool
+      f_spec = <f rhs> Int dNumInt |> co
+
+      RULE: forall d. f Int d = f_spec |> sym co
+
+  but the 'co' and 'sym co' are (a) playing no useful role, and (b) are
+  hard to generate.  At all costs we must avoid this:
+      RULE: forall d. f Int d |> co = f_spec
+  because the LHS will never match (indeed it's rejected in
+  decomposeRuleLhs).
+
+  So we simply do this:
+    - Generate a constraint to check that the specialised type (after
+      skolemiseation) is equal to the instantiated function type.
+    - But *discard* the evidence (coercion) for that constraint,
+      so that we ultimately generate the simpler code
+          f_spec :: Int -> F Int
+          f_spec = <f rhs> Int dNumInt
+
+          RULE: forall d. f Int d = f_spec
+      You can see this discarding happening in 
+
+3. Note that the HsWrapper can transform *any* function with the right
+   type prefix
+       forall ab. (Eq a, Ix b) => XXX
+   regardless of XXX.  It's sort of polymorphic in XXX.  This is
+   useful: we use the same wrapper to transform each of the class ops, as
+   well as the dict.  That's what goes on in TcInstDcls.mk_meth_spec_prags
+-}
+
 type PragFun = Name -> [LSig Name]
 
 mkPragFun :: [LSig Name] -> LHsBinds Name -> PragFun
@@ -879,7 +1010,7 @@ tcSpecPrags :: Id -> [LSig Name]
 tcSpecPrags poly_id prag_sigs
   = do { traceTc "tcSpecPrags" (ppr poly_id <+> ppr spec_sigs)
        ; unless (null bad_sigs) warn_discarded_sigs
-       ; pss <- mapAndRecoverM (wrapLocM (tcSpec poly_id)) spec_sigs
+       ; pss <- mapAndRecoverM (wrapLocM (tcSpecPrag poly_id)) spec_sigs
        ; return $ concatMap (\(L l ps) -> map (L l) ps) pss }
   where
     spec_sigs = filter isSpecLSig prag_sigs
@@ -891,20 +1022,24 @@ tcSpecPrags poly_id prag_sigs
 
 
 --------------
-tcSpec :: TcId -> Sig Name -> TcM [TcSpecPrag]
-tcSpec poly_id prag@(SpecSig fun_name hs_tys inl)
-  -- The Name fun_name in the SpecSig may not be the same as that of the poly_id
-  -- Example: SPECIALISE for a class method: the Name in the SpecSig is
-  --          for the selector Id, but the poly_id is something like $cop
-  -- However we want to use fun_name in the error message, since that is
-  -- what the user wrote (Trac #8537)
+tcSpecPrag :: TcId -> Sig Name -> TcM [TcSpecPrag]
+tcSpecPrag poly_id prag@(SpecSig fun_name hs_tys inl)
+-- See Note [Handling SPECIALISE pragmas]
+-- 
+-- The Name fun_name in the SpecSig may not be the same as that of the poly_id
+-- Example: SPECIALISE for a class method: the Name in the SpecSig is
+--          for the selector Id, but the poly_id is something like $cop
+-- However we want to use fun_name in the error message, since that is
+-- what the user wrote (Trac #8537)
   = addErrCtxt (spec_ctxt prag) $
     do  { spec_tys <- mapM (tcHsSigType sig_ctxt) hs_tys
         ; warnIf (not (isOverloadedTy poly_ty || isInlinePragma inl))
                  (ptext (sLit "SPECIALISE pragma for non-overloaded function")
                   <+> quotes (ppr fun_name))
                   -- Note [SPECIALISE pragmas]
-        ; wraps <- mapM (tcSubType sig_ctxt (idType poly_id)) spec_tys
+
+        ; wraps <- mapM (tcSpecWrapper sig_ctxt poly_ty) spec_tys
+        ; traceTc "tcSpecPrag" (ppr poly_id $$ nest 2 (vcat (map ppr (spec_tys `zip` wraps))))
         ; return [ (SpecPrag poly_id wrap inl) | wrap <- wraps ] }
   where
     name      = idName poly_id
@@ -912,7 +1047,24 @@ tcSpec poly_id prag@(SpecSig fun_name hs_tys inl)
     sig_ctxt  = FunSigCtxt name True
     spec_ctxt prag = hang (ptext (sLit "In the SPECIALISE pragma")) 2 (ppr prag)
 
-tcSpec _ prag = pprPanic "tcSpec" (ppr prag)
+tcSpecPrag _ prag = pprPanic "tcSpecPrag" (ppr prag)
+
+--------------
+tcSpecWrapper :: UserTypeCtxt -> TcType -> TcType -> TcM HsWrapper
+-- A simpler variant of tcSubType, used for SPECIALISE pragmas
+-- See Note [Handling SPECIALISE pragmas], wrinkle 1
+tcSpecWrapper ctxt poly_ty spec_ty 
+  = do { (sk_wrap, inst_wrap)
+               <- tcGen ctxt spec_ty $ \ _ spec_tau ->
+                  do { (inst_wrap, tau) <- deeplyInstantiate orig poly_ty
+                     ; _ <- unifyType spec_tau tau
+                            -- Deliberately ignore the evidence
+                            -- See Note [Handling SPECIALISE pragmas],
+                            --   wrinkle (2)
+                     ; return inst_wrap }
+       ; return (sk_wrap <.> inst_wrap) }
+  where
+    orig = SpecPragOrigin ctxt
 
 --------------
 tcImpPrags :: [LSig Name] -> TcM [LTcSpecPrag]
@@ -945,7 +1097,7 @@ tcImpSpec (name, prag)
  = do { id <- tcLookupId name
       ; unless (isAnyInlinePragma (idInlinePragma id))
                (addWarnTc (impSpecErr name))
-      ; tcSpec id prag }
+      ; tcSpecPrag id prag }
 
 impSpecErr :: Name -> SDoc
 impSpecErr name
@@ -957,7 +1109,13 @@ impSpecErr name
   where
     mod = nameModule name
 
---------------
+
+{- *********************************************************************
+*                                                                      *
+                         Vectorisation
+*                                                                      *
+********************************************************************* -}
+
 tcVectDecls :: [LVectDecl Name] -> TcM ([LVectDecl TcId])
 tcVectDecls decls
   = do { decls' <- mapM (wrapLocM tcVect) decls
@@ -1055,26 +1213,6 @@ vectCtxt thing = ptext (sLit "When checking the vectorisation declaration for")
 scalarTyConMustBeNullary :: MsgDoc
 scalarTyConMustBeNullary = ptext (sLit "VECTORISE SCALAR type constructor must be nullary")
 
---------------
--- If typechecking the binds fails, then return with each
--- signature-less binder given type (forall a.a), to minimise
--- subsequent error messages
-recoveryCode :: [Name] -> TcSigFun -> TcM (LHsBinds TcId, [Id], TopLevelFlag)
-recoveryCode binder_names sig_fn
-  = do  { traceTc "tcBindsWithSigs: error recovery" (ppr binder_names)
-        ; poly_ids <- mapM mk_dummy binder_names
-        ; return (emptyBag, poly_ids, if all is_closed poly_ids
-                                      then TopLevel else NotTopLevel) }
-  where
-    mk_dummy name
-        | isJust (sig_fn name) = tcLookupId name        -- Had signature; look it up
-        | otherwise            = return (mkLocalId name forall_a_a)    -- No signature
-
-    is_closed poly_id = isEmptyVarSet (tyVarsOfType (idType poly_id))
-
-forall_a_a :: TcType
-forall_a_a = mkForAllTy openAlphaTyVar (mkTyVarTy openAlphaTyVar)
-
 {-
 Note [SPECIALISE pragmas]
 ~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -1099,7 +1237,7 @@ for a non-overloaded function.
 
 ************************************************************************
 *                                                                      *
-\subsection{tcMonoBind}
+                         tcMonoBinds
 *                                                                      *
 ************************************************************************
 
@@ -1625,16 +1763,12 @@ strictBindErr flavour unlifted_bndrs binds
     msg | unlifted_bndrs = ptext (sLit "bindings for unlifted types")
         | otherwise      = ptext (sLit "bang-pattern or unboxed-tuple bindings")
 
-{-
-Note [Binding scoped type variables]
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-************************************************************************
+{- *********************************************************************
 *                                                                      *
-\subsection[TcBinds-errors]{Error contexts and messages}
+               Error contexts and messages
 *                                                                      *
-************************************************************************
--}
+********************************************************************* -}
 
 -- This one is called on LHS, when pat and grhss are both Name
 -- and on RHS, when pat is TcId and grhss is still Name

compiler/typecheck/TcInstDcls.hs

@@ -1566,6 +1566,7 @@ mk_meth_spec_prags :: Id -> [LTcSpecPrag] -> [LTcSpecPrag] -> TcSpecPrags
         --   * spec_prags_from_inst: derived from {-# SPECIALISE instance :: <blah> #-}
         --     These ones have the dfun inside, but [perhaps surprisingly]
         --     the correct wrapper.
+        -- See Note [Handling SPECIALISE pragmas] in TcBinds
 mk_meth_spec_prags meth_id spec_inst_prags spec_prags_for_me
   = SpecPrags (spec_prags_for_me ++ spec_prags_from_inst)
   where
@@ -1804,7 +1805,7 @@ tcSpecInst dfun_id prag@(SpecInstSig hs_ty)
   = addErrCtxt (spec_ctxt prag) $
     do  { (tyvars, theta, clas, tys) <- tcHsInstHead SpecInstCtxt hs_ty
         ; let spec_dfun_ty = mkDictFunTy tyvars theta clas tys
-        ; co_fn <- tcSubType SpecInstCtxt (idType dfun_id) spec_dfun_ty
+        ; co_fn <- tcSpecWrapper SpecInstCtxt (idType dfun_id) spec_dfun_ty
         ; return (SpecPrag dfun_id co_fn defaultInlinePragma) }
   where
     spec_ctxt prag = hang (ptext (sLit "In the SPECIALISE pragma")) 2 (ppr prag)

compiler/typecheck/TcRnTypes.hs

@@ -2033,10 +2033,11 @@ data CtOrigin
   = GivenOrigin SkolemInfo
 
   -- All the others are for *wanted* constraints
-  | OccurrenceOf Name           -- Occurrence of an overloaded identifier
-  | AppOrigin                   -- An application of some kind
+  | OccurrenceOf Name              -- Occurrence of an overloaded identifier
+  | AppOrigin                      -- An application of some kind
 
-  | SpecPragOrigin Name         -- Specialisation pragma for identifier
+  | SpecPragOrigin UserTypeCtxt    -- Specialisation pragma for 
+                                   -- function or instance
 
   | TypeEqOrigin { uo_actual   :: TcType
                  , uo_expected :: TcType }
@@ -2097,6 +2098,12 @@ pprCtOrigin :: CtOrigin -> SDoc
 
 pprCtOrigin (GivenOrigin sk) = ctoHerald <+> ppr sk
 
+pprCtOrigin (SpecPragOrigin ctxt) 
+  = case ctxt of
+       FunSigCtxt n _ -> ptext (sLit "a SPECIALISE pragma for") <+> quotes (ppr n)
+       SpecInstCtxt   -> ptext (sLit "a SPECIALISE INSTANCE pragma")
+       _              -> ptext (sLit "a SPECIALISE pragma")  -- Never happens I think
+
 pprCtOrigin (FunDepOrigin1 pred1 loc1 pred2 loc2)
   = hang (ctoHerald <+> ptext (sLit "a functional dependency between constraints:"))
        2 (vcat [ hang (quotes (ppr pred1)) 2 (pprArisingAt loc1)
@@ -2140,7 +2147,6 @@ pprCtOrigin simple_origin
 pprCtO :: CtOrigin -> SDoc  -- Ones that are short one-liners
 pprCtO (OccurrenceOf name)   = hsep [ptext (sLit "a use of"), quotes (ppr name)]
 pprCtO AppOrigin             = ptext (sLit "an application")
-pprCtO (SpecPragOrigin name) = hsep [ptext (sLit "a specialisation pragma for"), quotes (ppr name)]
 pprCtO (IPOccOrigin name)    = hsep [ptext (sLit "a use of implicit parameter"), quotes (ppr name)]
 pprCtO RecordUpdOrigin       = ptext (sLit "a record update")
 pprCtO ExprSigOrigin         = ptext (sLit "an expression type signature")

testsuite/tests/parser/should_fail/T7848.stderr

@@ -22,18 +22,15 @@ T7848.hs:6:57:
               (&) = x
 
 T7848.hs:10:9:
-    Couldn't match type ‘a’ with ‘t -> t1 -> A -> A -> A -> A -> t2’
+    Couldn't match expected type ‘t -> t1 -> A -> A -> A -> A -> t2’
+                with actual type ‘a’
       ‘a’ is a rigid type variable bound by
           the type signature for: (&) :: a at T7848.hs:10:9
-    Expected type: forall a. a
-      Actual type: t -> t1 -> A -> A -> A -> A -> t2
     Relevant bindings include
       z :: t1 (bound at T7848.hs:6:12)
       (&) :: t1 (bound at T7848.hs:6:8)
       (+) :: t (bound at T7848.hs:6:3)
       x :: t -> t1 -> A -> A -> A -> A -> t2 (bound at T7848.hs:6:1)
-    When checking that: t -> t1 -> A -> A -> A -> A -> t2
-      is more polymorphic than: forall a. a
     In the SPECIALISE pragma {-# SPECIALIZE (&) :: a #-}
     In an equation for ‘x’:
         x (+) ((&)@z) ((:&&) a b) (c :&& d) (e `A` f) (A g h)