+tcSpecPrags poly_id prag_sigs
+ = do { unless (null bad_sigs) warn_discarded_sigs
+ ; mapAndRecoverM (wrapLocM (tcSpec poly_id)) spec_sigs }
+ where
+ spec_sigs = filter isSpecLSig prag_sigs
+ bad_sigs = filter is_bad_sig prag_sigs
+ is_bad_sig s = not (isSpecLSig s || isInlineLSig s)
+
+ warn_discarded_sigs = warnPrags poly_id bad_sigs $
+ ptext (sLit "Discarding unexpected pragmas for")
+
+
+--------------
+tcSpec :: TcId -> Sig Name -> TcM TcSpecPrag
+tcSpec poly_id prag@(SpecSig _ hs_ty inl)
+ -- The 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
+ = addErrCtxt (spec_ctxt prag) $
+ do { spec_ty <- tcHsSigType sig_ctxt hs_ty
+ ; warnIf (not (isOverloadedTy poly_ty || isInlinePragma inl))
+ (ptext (sLit "SPECIALISE pragma for non-overloaded function") <+> quotes (ppr poly_id))
+ -- Note [SPECIALISE pragmas]
+ ; wrap <- tcSubType origin sig_ctxt (idType poly_id) spec_ty
+ ; return (SpecPrag poly_id wrap inl) }
+ where
+ name = idName poly_id
+ poly_ty = idType poly_id
+ origin = SpecPragOrigin name
+ sig_ctxt = FunSigCtxt name
+ spec_ctxt prag = hang (ptext (sLit "In the SPECIALISE pragma")) 2 (ppr prag)
+
+tcSpec _ prag = pprPanic "tcSpec" (ppr prag)
+
+--------------
+tcImpPrags :: [LSig Name] -> TcM [LTcSpecPrag]
+-- SPECIALISE pragamas for imported things
+tcImpPrags prags
+ = do { this_mod <- getModule
+ ; dflags <- getDOpts
+ ; if (not_specialising dflags) then
+ return []
+ else
+ mapAndRecoverM (wrapLocM tcImpSpec)
+ [L loc (name,prag) | (L loc prag@(SpecSig (L _ name) _ _)) <- prags
+ , not (nameIsLocalOrFrom this_mod name) ] }
+ where
+ -- Ignore SPECIALISE pragmas for imported things
+ -- when we aren't specialising, or when we aren't generating
+ -- code. The latter happens when Haddocking the base library;
+ -- we don't wnat complaints about lack of INLINABLE pragmas
+ not_specialising dflags
+ | not (dopt Opt_Specialise dflags) = True
+ | otherwise = case hscTarget dflags of
+ HscNothing -> True
+ HscInterpreted -> True
+ _other -> False
+
+tcImpSpec :: (Name, Sig Name) -> TcM TcSpecPrag
+tcImpSpec (name, prag)
+ = do { id <- tcLookupId name
+ ; unless (isAnyInlinePragma (idInlinePragma id))
+ (addWarnTc (impSpecErr name))
+ ; tcSpec id prag }
+
+impSpecErr :: Name -> SDoc
+impSpecErr name
+ = hang (ptext (sLit "You cannot SPECIALISE") <+> quotes (ppr name))
+ 2 (vcat [ ptext (sLit "because its definition has no INLINE/INLINABLE pragma")
+ , parens $ sep
+ [ ptext (sLit "or its defining module") <+> quotes (ppr mod)
+ , ptext (sLit "was compiled without -O")]])
+ where
+ mod = nameModule name
+
+--------------
+tcVectDecls :: [LVectDecl Name] -> TcM ([LVectDecl TcId])
+tcVectDecls decls
+ = do { decls' <- mapM (wrapLocM tcVect) decls
+ ; let ids = map lvectDeclName decls'
+ dups = findDupsEq (==) ids
+ ; mapM_ reportVectDups dups
+ ; traceTcConstraints "End of tcVectDecls"
+ ; return decls'
+ }
+ where
+ reportVectDups (first:_second:_more)
+ = addErrAt (getSrcSpan first) $
+ ptext (sLit "Duplicate vectorisation declarations for") <+> ppr first
+ reportVectDups _ = return ()
+
+--------------
+tcVect :: VectDecl Name -> TcM (VectDecl TcId)
+-- We can't typecheck the expression of a vectorisation declaration against the vectorised type
+-- of the original definition as this requires internals of the vectoriser not available during
+-- type checking. Instead, we infer the type of the expression and leave it to the vectoriser
+-- to check the compatibility of the Core types.
+tcVect (HsVect name Nothing)
+ = addErrCtxt (vectCtxt name) $
+ do { id <- wrapLocM tcLookupId name
+ ; return $ HsVect id Nothing
+ }
+tcVect (HsVect name@(L loc _) (Just rhs))
+ = addErrCtxt (vectCtxt name) $
+ do { _id <- wrapLocM tcLookupId name -- need to ensure that the name is already defined
+
+ -- turn the vectorisation declaration into a single non-recursive binding
+ ; let bind = L loc $ mkFunBind name [mkSimpleMatch [] rhs]
+ sigFun = const Nothing
+ pragFun = mkPragFun [] (unitBag bind)
+
+ -- perform type inference (including generalisation)
+ ; (binds, [id']) <- tcPolyInfer TopLevel False sigFun pragFun NonRecursive [bind]
+
+ ; traceTc "tcVect inferred type" $ ppr (varType id')
+ ; traceTc "tcVect bindings" $ ppr binds
+
+ -- add all bindings, including the type variable and dictionary bindings produced by type
+ -- generalisation to the right-hand side of the vectorisation declaration
+ ; let [AbsBinds tvs evs _ evBinds actualBinds] = (map unLoc . bagToList) binds
+ ; let [bind'] = bagToList actualBinds
+ MatchGroup
+ [L _ (Match _ _ (GRHSs [L _ (GRHS _ rhs')] _))]
+ _ = (fun_matches . unLoc) bind'
+ rhsWrapped = mkHsLams tvs evs (mkHsDictLet evBinds rhs')
+
+ -- We return the type-checked 'Id', to propagate the inferred signature
+ -- to the vectoriser - see "Note [Typechecked vectorisation pragmas]" in HsDecls
+ ; return $ HsVect (L loc id') (Just rhsWrapped)
+ }
+tcVect (HsNoVect name)
+ = addErrCtxt (vectCtxt name) $
+ do { id <- wrapLocM tcLookupId name
+ ; return $ HsNoVect id
+ }
+
+vectCtxt :: Located Name -> SDoc
+vectCtxt name = ptext (sLit "When checking the vectorisation declaration for") <+> ppr name
+