--- lcl_id -> (gbl_id, prags)
-mkABEnv exports = mkVarEnv [ (lcl_id, (gbl_id, prags))
- | (_, gbl_id, lcl_id, prags) <- exports]
-
-
-dsSpec :: [TyVar] -> [DictId] -> [TyVar]
- -> Id -> Id -- Global, local
- -> CoreBind -> Prag
- -> DsM (Maybe ((Id,CoreExpr), -- Binding for specialised Id
- CoreRule)) -- Rule for the Global Id
-
--- Example:
--- f :: (Eq a, Ix b) => a -> b -> b
--- {-# SPECIALISE f :: Ix b => Int -> b -> b #-}
---
--- AbsBinds [ab] [d1,d2] [([ab], f, f_mono, prags)] binds
---
--- SpecPrag (/\b.\(d:Ix b). f Int b dInt d)
--- (forall b. Ix b => Int -> b -> b)
---
--- Rule: forall b,(d:Ix b). f Int b dInt d = f_spec b d
---
--- Spec bind: f_spec = Let f = /\ab \(d1:Eq a)(d2:Ix b). let binds in f_mono
--- /\b.\(d:Ix b). in f Int b dInt d
--- The idea is that f occurs just once, so it'll be
--- inlined and specialised
-
-dsSpec all_tvs dicts tvs poly_id mono_id mono_bind (InlinePrag {})
- = return Nothing
-
-dsSpec all_tvs dicts tvs poly_id mono_id mono_bind
- (SpecPrag spec_expr spec_ty const_dicts inl)
- = do { let poly_name = idName poly_id
- ; spec_name <- newLocalName poly_name
- ; ds_spec_expr <- dsExpr spec_expr
- ; let (bndrs, body) = collectBinders ds_spec_expr
- mb_lhs = decomposeRuleLhs (bndrs ++ const_dicts) body
-
- ; case mb_lhs of
- Nothing -> do { warnDs msg; return Nothing }
-
- Just (bndrs', var, args) -> return (Just (addInlineInfo inl spec_id spec_rhs, rule))
- where
- local_poly = setIdNotExported poly_id
- -- Very important to make the 'f' non-exported,
- -- else it won't be inlined!
- spec_id = mkLocalId spec_name spec_ty
- spec_rhs = Let (NonRec local_poly poly_f_body) ds_spec_expr
- poly_f_body = mkLams (tvs ++ dicts) $
- fix_up (Let mono_bind (Var mono_id))
-
- -- Quantify over constant dicts on the LHS, since
- -- their value depends only on their type
- -- The ones we are interested in may even be imported
- -- e.g. GHC.Base.dEqInt
-
- rule = mkLocalRule (mkFastString ("SPEC " ++ showSDoc (ppr poly_name)))
- AlwaysActive poly_name
- bndrs' -- Includes constant dicts
- args
- (mkVarApps (Var spec_id) bndrs)
- }
- where
+-- lcl_id -> (tyvars, gbl_id, lcl_id, prags)
+mkABEnv exports = mkVarEnv [ (lcl_id, export) | export@(_, _, lcl_id, _) <- exports]
+
+mkArityEnv :: LHsBinds Id -> IdEnv Arity
+ -- Maps a local to the arity of its definition
+mkArityEnv binds = foldrBag (plusVarEnv . lhsBindArity) emptyVarEnv binds
+
+lhsBindArity :: LHsBind Id -> IdEnv Arity
+lhsBindArity (L _ (FunBind { fun_id = id, fun_matches = ms }))
+ = unitVarEnv (unLoc id) (matchGroupArity ms)
+lhsBindArity (L _ (AbsBinds { abs_exports = exports
+ , abs_dicts = dicts
+ , abs_binds = binds }))
+ = mkVarEnv [ (gbl, lookupArity ar_env lcl + n_val_dicts)
+ | (_, gbl, lcl, _) <- exports]
+ where -- See Note [Nested arities]
+ ar_env = mkArityEnv binds
+ n_val_dicts = dictArity dicts
+
+lhsBindArity _ = emptyVarEnv -- PatBind/VarBind
+
+dictArity :: [Var] -> Arity
+-- Don't count coercion variables in arity
+dictArity dicts = count isId dicts
+
+lookupArity :: IdEnv Arity -> Id -> Arity
+lookupArity ar_env id = lookupVarEnv ar_env id `orElse` 0
+\end{code}
+
+Note [Eta-expanding INLINE things]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Consider
+ foo :: Eq a => a -> a
+ {-# INLINE foo #-}
+ foo x = ...
+
+If (foo d) ever gets floated out as a common sub-expression (which can
+happen as a result of method sharing), there's a danger that we never
+get to do the inlining, which is a Terribly Bad thing given that the
+user said "inline"!
+
+To avoid this we pre-emptively eta-expand the definition, so that foo
+has the arity with which it is declared in the source code. In this
+example it has arity 2 (one for the Eq and one for x). Doing this
+should mean that (foo d) is a PAP and we don't share it.
+
+Note [Nested arities]
+~~~~~~~~~~~~~~~~~~~~~
+For reasons that are not entirely clear, method bindings come out looking like
+this:
+
+ AbsBinds [] [] [$cfromT <= [] fromT]
+ $cfromT [InlPrag=INLINE] :: T Bool -> Bool
+ { AbsBinds [] [] [fromT <= [] fromT_1]
+ fromT :: T Bool -> Bool
+ { fromT_1 ((TBool b)) = not b } } }
+
+Note the nested AbsBind. The arity for the InlineRule on $cfromT should be
+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 #-}
+
+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) => <blah>
+regardless of <blah>. It's sort of polymorphic in <blah>. 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 (/\ab \d1 d2. Let binds in f_mono)
+
+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 has a *copy* of 'binds', so that it can fully
+ specialise it.
+
+\begin{code}
+------------------------
+dsSpecs :: [TyVar] -> [DictId] -> [TyVar]
+ -> Id -> Id -> Arity -- Global, local, arity of local
+ -> CoreBind -> [LSpecPrag]
+ -> DsM ( [(Id,CoreExpr)] -- Binding for specialised Ids
+ , [CoreRule] ) -- Rules for the Global Ids
+-- See Note [Implementing SPECIALISE pragmas]
+dsSpecs all_tvs dicts tvs poly_id mono_id inl_arity mono_bind prags
+ = do { pairs <- mapMaybeM spec_one prags
+ ; let (spec_binds_s, rules) = unzip pairs
+ ; return (concat spec_binds_s, rules) }
+ where
+ spec_one :: LSpecPrag -> DsM (Maybe ([(Id,CoreExpr)], CoreRule))
+ spec_one (L loc (SpecPrag spec_co spec_inl))
+ = putSrcSpanDs loc $
+ do { let poly_name = idName poly_id
+ ; spec_name <- newLocalName poly_name
+ ; wrap_fn <- dsCoercion spec_co
+ ; let ds_spec_expr = wrap_fn (Var poly_id)
+ ; case decomposeRuleLhs ds_spec_expr of {
+ Nothing -> do { warnDs (decomp_msg spec_co)
+ ; return Nothing } ;
+
+ Just (bndrs, _fn, args) ->
+
+ -- Check for dead binders: Note [Unused spec binders]
+ case filter isDeadBinder bndrs of {
+ bs | not (null bs) -> do { warnDs (dead_msg bs); return Nothing }
+ | otherwise -> do
+
+ { (spec_unf, unf_pairs) <- specUnfolding wrap_fn (realIdUnfolding poly_id)
+
+ ; let f_body = fix_up (Let mono_bind (Var mono_id))
+ spec_ty = exprType ds_spec_expr
+ spec_id = mkLocalId spec_name spec_ty
+ `setInlinePragma` inl_prag
+ `setIdUnfolding` spec_unf
+ inl_prag | isDefaultInlinePragma spec_inl = idInlinePragma poly_id
+ | otherwise = spec_inl
+ -- Get the INLINE pragma from SPECIALISE declaration, or,
+ -- failing that, from the original Id
+
+ spec_id_arity = inl_arity + count isDictId bndrs
+
+ extra_dict_bndrs = [ localiseId d -- See Note [Constant rule dicts]
+ | d <- varSetElems (exprFreeVars ds_spec_expr)
+ , isDictId d]
+ -- Note [Const rule dicts]
+
+ rule = mkLocalRule (mkFastString ("SPEC " ++ showSDoc (ppr poly_name)))
+ AlwaysActive poly_name
+ (extra_dict_bndrs ++ bndrs) args
+ (mkVarApps (Var spec_id) bndrs)
+
+ spec_rhs = wrap_fn (mkLams (tvs ++ dicts) f_body)
+ spec_pair = makeCorePair spec_id spec_id_arity spec_rhs
+
+ ; return (Just (spec_pair : unf_pairs, rule))
+ } } } }
+