+ prs = mapCatMaybes get_sig sigs
+
+ get_sig :: LSig Name -> Maybe (Located Name, LSig Name)
+ get_sig (L l (SpecSig nm ty inl)) = Just (nm, L l $ SpecSig nm ty (add_arity nm inl))
+ get_sig (L l (InlineSig nm inl)) = Just (nm, L l $ InlineSig nm (add_arity nm inl))
+ get_sig _ = Nothing
+
+ add_arity (L _ n) inl_prag -- Adjust inl_sat field to match visible arity of function
+ | Just ar <- lookupNameEnv ar_env n,
+ Inline <- inl_inline inl_prag = inl_prag { inl_sat = Just ar }
+ -- add arity only for real INLINE pragmas, not INLINABLE
+ | otherwise = inl_prag
+
+ prag_env :: NameEnv [LSig Name]
+ prag_env = foldl add emptyNameEnv prs
+ add env (L _ n,p) = extendNameEnv_Acc (:) singleton env n p
+
+ -- ar_env maps a local to the arity of its definition
+ ar_env :: NameEnv Arity
+ ar_env = foldrBag lhsBindArity emptyNameEnv binds
+
+lhsBindArity :: LHsBind Name -> NameEnv Arity -> NameEnv Arity
+lhsBindArity (L _ (FunBind { fun_id = id, fun_matches = ms })) env
+ = extendNameEnv env (unLoc id) (matchGroupArity ms)
+lhsBindArity _ env = env -- PatBind/VarBind
+
+------------------
+tcSpecPrags :: Id -> [LSig Name]
+ -> TcM [LTcSpecPrag]
+-- Add INLINE and SPECIALSE pragmas
+-- INLINE prags are added to the (polymorphic) Id directly
+-- SPECIALISE prags are passed to the desugarer via TcSpecPrags
+-- Pre-condition: the poly_id is zonked
+-- Reason: required by tcSubExp
+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]
+tcImpPrags prags
+ = do { this_mod <- getModule
+ ; mapAndRecoverM (wrapLocM tcImpSpec)
+ [L loc (name,prag) | (L loc prag@(SpecSig (L _ name) _ _)) <- prags
+ , not (nameIsLocalOrFrom this_mod name) ] }
+
+tcImpSpec :: (Name, Sig Name) -> TcM TcSpecPrag
+tcImpSpec (name, prag)
+ = do { id <- tcLookupId name
+ ; checkTc (isAnyInlinePragma (idInlinePragma id))
+ (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")
+ , ptext (sLit "(or you compiled its defining module without -O)")])
+
+--------------
+tcVectDecls :: [LVectDecl Name] -> TcM ([LVectDecl TcId])
+tcVectDecls decls
+ = do { decls' <- mapM (wrapLocM tcVect) decls
+ ; let ids = [unLoc id | L _ (HsVect id _) <- 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)
+ }
+
+vectCtxt :: Located Name -> SDoc
+vectCtxt name = ptext (sLit "When checking the vectorisation declaration for") <+> ppr name
+