PageRenderTime 66ms CodeModel.GetById 22ms RepoModel.GetById 1ms app.codeStats 0ms

/ghc-7.0.4/compiler/specialise/SpecConstr.lhs

http://picorec.googlecode.com/
Haskell | 1680 lines | 1131 code | 309 blank | 240 comment | 54 complexity | 666b76979e6589c961b23b12ee209171 MD5 | raw file
Possible License(s): BSD-3-Clause, BSD-2-Clause

Large files files are truncated, but you can click here to view the full file

  1. %
  2. % (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
  3. %
  4. \section[SpecConstr]{Specialise over constructors}
  5. \begin{code}
  6. -- The above warning supression flag is a temporary kludge.
  7. -- While working on this module you are encouraged to remove it and fix
  8. -- any warnings in the module. See
  9. -- http://hackage.haskell.org/trac/ghc/wiki/Commentary/CodingStyle#Warnings
  10. -- for details
  11. module SpecConstr(
  12. specConstrProgram
  13. #ifdef GHCI
  14. , SpecConstrAnnotation(..)
  15. #endif
  16. ) where
  17. #include "HsVersions.h"
  18. import CoreSyn
  19. import CoreSubst
  20. import CoreUtils
  21. import CoreUnfold ( couldBeSmallEnoughToInline )
  22. import CoreFVs ( exprsFreeVars )
  23. import CoreMonad
  24. import HscTypes ( ModGuts(..) )
  25. import WwLib ( mkWorkerArgs )
  26. import DataCon
  27. import Coercion
  28. import Rules
  29. import Type hiding( substTy )
  30. import Id
  31. import MkCore ( mkImpossibleExpr )
  32. import Var
  33. import VarEnv
  34. import VarSet
  35. import Name
  36. import BasicTypes
  37. import DynFlags ( DynFlags(..) )
  38. import StaticFlags ( opt_PprStyle_Debug )
  39. import Maybes ( orElse, catMaybes, isJust, isNothing )
  40. import Demand
  41. import DmdAnal ( both )
  42. import Serialized ( deserializeWithData )
  43. import Util
  44. import UniqSupply
  45. import Outputable
  46. import FastString
  47. import UniqFM
  48. import MonadUtils
  49. import Control.Monad ( zipWithM )
  50. import Data.List
  51. -- See Note [SpecConstrAnnotation]
  52. #ifndef GHCI
  53. type SpecConstrAnnotation = ()
  54. #else
  55. import Literal ( literalType )
  56. import TyCon ( TyCon )
  57. import GHC.Exts( SpecConstrAnnotation(..) )
  58. #endif
  59. \end{code}
  60. -----------------------------------------------------
  61. Game plan
  62. -----------------------------------------------------
  63. Consider
  64. drop n [] = []
  65. drop 0 xs = []
  66. drop n (x:xs) = drop (n-1) xs
  67. After the first time round, we could pass n unboxed. This happens in
  68. numerical code too. Here's what it looks like in Core:
  69. drop n xs = case xs of
  70. [] -> []
  71. (y:ys) -> case n of
  72. I# n# -> case n# of
  73. 0 -> []
  74. _ -> drop (I# (n# -# 1#)) xs
  75. Notice that the recursive call has an explicit constructor as argument.
  76. Noticing this, we can make a specialised version of drop
  77. RULE: drop (I# n#) xs ==> drop' n# xs
  78. drop' n# xs = let n = I# n# in ...orig RHS...
  79. Now the simplifier will apply the specialisation in the rhs of drop', giving
  80. drop' n# xs = case xs of
  81. [] -> []
  82. (y:ys) -> case n# of
  83. 0 -> []
  84. _ -> drop (n# -# 1#) xs
  85. Much better!
  86. We'd also like to catch cases where a parameter is carried along unchanged,
  87. but evaluated each time round the loop:
  88. f i n = if i>0 || i>n then i else f (i*2) n
  89. Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  90. In Core, by the time we've w/wd (f is strict in i) we get
  91. f i# n = case i# ># 0 of
  92. False -> I# i#
  93. True -> case n of n' { I# n# ->
  94. case i# ># n# of
  95. False -> I# i#
  96. True -> f (i# *# 2#) n'
  97. At the call to f, we see that the argument, n is know to be (I# n#),
  98. and n is evaluated elsewhere in the body of f, so we can play the same
  99. trick as above.
  100. Note [Reboxing]
  101. ~~~~~~~~~~~~~~~
  102. We must be careful not to allocate the same constructor twice. Consider
  103. f p = (...(case p of (a,b) -> e)...p...,
  104. ...let t = (r,s) in ...t...(f t)...)
  105. At the recursive call to f, we can see that t is a pair. But we do NOT want
  106. to make a specialised copy:
  107. f' a b = let p = (a,b) in (..., ...)
  108. because now t is allocated by the caller, then r and s are passed to the
  109. recursive call, which allocates the (r,s) pair again.
  110. This happens if
  111. (a) the argument p is used in other than a case-scrutinsation way.
  112. (b) the argument to the call is not a 'fresh' tuple; you have to
  113. look into its unfolding to see that it's a tuple
  114. Hence the "OR" part of Note [Good arguments] below.
  115. ALTERNATIVE 2: pass both boxed and unboxed versions. This no longer saves
  116. allocation, but does perhaps save evals. In the RULE we'd have
  117. something like
  118. f (I# x#) = f' (I# x#) x#
  119. If at the call site the (I# x) was an unfolding, then we'd have to
  120. rely on CSE to eliminate the duplicate allocation.... This alternative
  121. doesn't look attractive enough to pursue.
  122. ALTERNATIVE 3: ignore the reboxing problem. The trouble is that
  123. the conservative reboxing story prevents many useful functions from being
  124. specialised. Example:
  125. foo :: Maybe Int -> Int -> Int
  126. foo (Just m) 0 = 0
  127. foo x@(Just m) n = foo x (n-m)
  128. Here the use of 'x' will clearly not require boxing in the specialised function.
  129. The strictness analyser has the same problem, in fact. Example:
  130. f p@(a,b) = ...
  131. If we pass just 'a' and 'b' to the worker, it might need to rebox the
  132. pair to create (a,b). A more sophisticated analysis might figure out
  133. precisely the cases in which this could happen, but the strictness
  134. analyser does no such analysis; it just passes 'a' and 'b', and hopes
  135. for the best.
  136. So my current choice is to make SpecConstr similarly aggressive, and
  137. ignore the bad potential of reboxing.
  138. Note [Good arguments]
  139. ~~~~~~~~~~~~~~~~~~~~~
  140. So we look for
  141. * A self-recursive function. Ignore mutual recursion for now,
  142. because it's less common, and the code is simpler for self-recursion.
  143. * EITHER
  144. a) At a recursive call, one or more parameters is an explicit
  145. constructor application
  146. AND
  147. That same parameter is scrutinised by a case somewhere in
  148. the RHS of the function
  149. OR
  150. b) At a recursive call, one or more parameters has an unfolding
  151. that is an explicit constructor application
  152. AND
  153. That same parameter is scrutinised by a case somewhere in
  154. the RHS of the function
  155. AND
  156. Those are the only uses of the parameter (see Note [Reboxing])
  157. What to abstract over
  158. ~~~~~~~~~~~~~~~~~~~~~
  159. There's a bit of a complication with type arguments. If the call
  160. site looks like
  161. f p = ...f ((:) [a] x xs)...
  162. then our specialised function look like
  163. f_spec x xs = let p = (:) [a] x xs in ....as before....
  164. This only makes sense if either
  165. a) the type variable 'a' is in scope at the top of f, or
  166. b) the type variable 'a' is an argument to f (and hence fs)
  167. Actually, (a) may hold for value arguments too, in which case
  168. we may not want to pass them. Supose 'x' is in scope at f's
  169. defn, but xs is not. Then we'd like
  170. f_spec xs = let p = (:) [a] x xs in ....as before....
  171. Similarly (b) may hold too. If x is already an argument at the
  172. call, no need to pass it again.
  173. Finally, if 'a' is not in scope at the call site, we could abstract
  174. it as we do the term variables:
  175. f_spec a x xs = let p = (:) [a] x xs in ...as before...
  176. So the grand plan is:
  177. * abstract the call site to a constructor-only pattern
  178. e.g. C x (D (f p) (g q)) ==> C s1 (D s2 s3)
  179. * Find the free variables of the abstracted pattern
  180. * Pass these variables, less any that are in scope at
  181. the fn defn. But see Note [Shadowing] below.
  182. NOTICE that we only abstract over variables that are not in scope,
  183. so we're in no danger of shadowing variables used in "higher up"
  184. in f_spec's RHS.
  185. Note [Shadowing]
  186. ~~~~~~~~~~~~~~~~
  187. In this pass we gather up usage information that may mention variables
  188. that are bound between the usage site and the definition site; or (more
  189. seriously) may be bound to something different at the definition site.
  190. For example:
  191. f x = letrec g y v = let x = ...
  192. in ...(g (a,b) x)...
  193. Since 'x' is in scope at the call site, we may make a rewrite rule that
  194. looks like
  195. RULE forall a,b. g (a,b) x = ...
  196. But this rule will never match, because it's really a different 'x' at
  197. the call site -- and that difference will be manifest by the time the
  198. simplifier gets to it. [A worry: the simplifier doesn't *guarantee*
  199. no-shadowing, so perhaps it may not be distinct?]
  200. Anyway, the rule isn't actually wrong, it's just not useful. One possibility
  201. is to run deShadowBinds before running SpecConstr, but instead we run the
  202. simplifier. That gives the simplest possible program for SpecConstr to
  203. chew on; and it virtually guarantees no shadowing.
  204. Note [Specialising for constant parameters]
  205. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  206. This one is about specialising on a *constant* (but not necessarily
  207. constructor) argument
  208. foo :: Int -> (Int -> Int) -> Int
  209. foo 0 f = 0
  210. foo m f = foo (f m) (+1)
  211. It produces
  212. lvl_rmV :: GHC.Base.Int -> GHC.Base.Int
  213. lvl_rmV =
  214. \ (ds_dlk :: GHC.Base.Int) ->
  215. case ds_dlk of wild_alH { GHC.Base.I# x_alG ->
  216. GHC.Base.I# (GHC.Prim.+# x_alG 1)
  217. T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
  218. GHC.Prim.Int#
  219. T.$wfoo =
  220. \ (ww_sme :: GHC.Prim.Int#) (w_smg :: GHC.Base.Int -> GHC.Base.Int) ->
  221. case ww_sme of ds_Xlw {
  222. __DEFAULT ->
  223. case w_smg (GHC.Base.I# ds_Xlw) of w1_Xmo { GHC.Base.I# ww1_Xmz ->
  224. T.$wfoo ww1_Xmz lvl_rmV
  225. };
  226. 0 -> 0
  227. }
  228. The recursive call has lvl_rmV as its argument, so we could create a specialised copy
  229. with that argument baked in; that is, not passed at all. Now it can perhaps be inlined.
  230. When is this worth it? Call the constant 'lvl'
  231. - If 'lvl' has an unfolding that is a constructor, see if the corresponding
  232. parameter is scrutinised anywhere in the body.
  233. - If 'lvl' has an unfolding that is a inlinable function, see if the corresponding
  234. parameter is applied (...to enough arguments...?)
  235. Also do this is if the function has RULES?
  236. Also
  237. Note [Specialising for lambda parameters]
  238. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  239. foo :: Int -> (Int -> Int) -> Int
  240. foo 0 f = 0
  241. foo m f = foo (f m) (\n -> n-m)
  242. This is subtly different from the previous one in that we get an
  243. explicit lambda as the argument:
  244. T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
  245. GHC.Prim.Int#
  246. T.$wfoo =
  247. \ (ww_sm8 :: GHC.Prim.Int#) (w_sma :: GHC.Base.Int -> GHC.Base.Int) ->
  248. case ww_sm8 of ds_Xlr {
  249. __DEFAULT ->
  250. case w_sma (GHC.Base.I# ds_Xlr) of w1_Xmf { GHC.Base.I# ww1_Xmq ->
  251. T.$wfoo
  252. ww1_Xmq
  253. (\ (n_ad3 :: GHC.Base.Int) ->
  254. case n_ad3 of wild_alB { GHC.Base.I# x_alA ->
  255. GHC.Base.I# (GHC.Prim.-# x_alA ds_Xlr)
  256. })
  257. };
  258. 0 -> 0
  259. }
  260. I wonder if SpecConstr couldn't be extended to handle this? After all,
  261. lambda is a sort of constructor for functions and perhaps it already
  262. has most of the necessary machinery?
  263. Furthermore, there's an immediate win, because you don't need to allocate the lamda
  264. at the call site; and if perchance it's called in the recursive call, then you
  265. may avoid allocating it altogether. Just like for constructors.
  266. Looks cool, but probably rare...but it might be easy to implement.
  267. Note [SpecConstr for casts]
  268. ~~~~~~~~~~~~~~~~~~~~~~~~~~~
  269. Consider
  270. data family T a :: *
  271. data instance T Int = T Int
  272. foo n = ...
  273. where
  274. go (T 0) = 0
  275. go (T n) = go (T (n-1))
  276. The recursive call ends up looking like
  277. go (T (I# ...) `cast` g)
  278. So we want to spot the construtor application inside the cast.
  279. That's why we have the Cast case in argToPat
  280. Note [Local recursive groups]
  281. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  282. For a *local* recursive group, we can see all the calls to the
  283. function, so we seed the specialisation loop from the calls in the
  284. body, not from the calls in the RHS. Consider:
  285. bar m n = foo n (n,n) (n,n) (n,n) (n,n)
  286. where
  287. foo n p q r s
  288. | n == 0 = m
  289. | n > 3000 = case p of { (p1,p2) -> foo (n-1) (p2,p1) q r s }
  290. | n > 2000 = case q of { (q1,q2) -> foo (n-1) p (q2,q1) r s }
  291. | n > 1000 = case r of { (r1,r2) -> foo (n-1) p q (r2,r1) s }
  292. | otherwise = case s of { (s1,s2) -> foo (n-1) p q r (s2,s1) }
  293. If we start with the RHSs of 'foo', we get lots and lots of specialisations,
  294. most of which are not needed. But if we start with the (single) call
  295. in the rhs of 'bar' we get exactly one fully-specialised copy, and all
  296. the recursive calls go to this fully-specialised copy. Indeed, the original
  297. function is later collected as dead code. This is very important in
  298. specialising the loops arising from stream fusion, for example in NDP where
  299. we were getting literally hundreds of (mostly unused) specialisations of
  300. a local function.
  301. Note [Do not specialise diverging functions]
  302. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  303. Specialising a function that just diverges is a waste of code.
  304. Furthermore, it broke GHC (simpl014) thus:
  305. {-# STR Sb #-}
  306. f = \x. case x of (a,b) -> f x
  307. If we specialise f we get
  308. f = \x. case x of (a,b) -> fspec a b
  309. But fspec doesn't have decent strictnes info. As it happened,
  310. (f x) :: IO t, so the state hack applied and we eta expanded fspec,
  311. and hence f. But now f's strictness is less than its arity, which
  312. breaks an invariant.
  313. Note [SpecConstrAnnotation]
  314. ~~~~~~~~~~~~~~~~~~~~~~~~~~~
  315. SpecConstrAnnotation is defined in GHC.Exts, and is only guaranteed to
  316. be available in stage 2 (well, until the bootstrap compiler can be
  317. guaranteed to have it)
  318. So we define it to be () in stage1 (ie when GHCI is undefined), and
  319. '#ifdef' out the code that uses it.
  320. See also Note [Forcing specialisation]
  321. Note [Forcing specialisation]
  322. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  323. With stream fusion and in other similar cases, we want to fully specialise
  324. some (but not necessarily all!) loops regardless of their size and the
  325. number of specialisations. We allow a library to specify this by annotating
  326. a type with ForceSpecConstr and then adding a parameter of that type to the
  327. loop. Here is a (simplified) example from the vector library:
  328. data SPEC = SPEC | SPEC2
  329. {-# ANN type SPEC ForceSpecConstr #-}
  330. foldl :: (a -> b -> a) -> a -> Stream b -> a
  331. {-# INLINE foldl #-}
  332. foldl f z (Stream step s _) = foldl_loop SPEC z s
  333. where
  334. foldl_loop SPEC z s = case step s of
  335. Yield x s' -> foldl_loop SPEC (f z x) s'
  336. Skip -> foldl_loop SPEC z s'
  337. Done -> z
  338. SpecConstr will spot the SPEC parameter and always fully specialise
  339. foldl_loop. Note that we can't just annotate foldl_loop since it isn't a
  340. top-level function but even if we could, inlining etc. could easily drop the
  341. annotation. We also have to prevent the SPEC argument from being removed by
  342. w/w which is why SPEC is a sum type. This is all quite ugly; we ought to come
  343. up with a better design.
  344. ForceSpecConstr arguments are spotted in scExpr' and scTopBinds which then set
  345. force_spec to True when calling specLoop. This flag makes specLoop and
  346. specialise ignore specConstrCount and specConstrThreshold when deciding
  347. whether to specialise a function.
  348. -----------------------------------------------------
  349. Stuff not yet handled
  350. -----------------------------------------------------
  351. Here are notes arising from Roman's work that I don't want to lose.
  352. Example 1
  353. ~~~~~~~~~
  354. data T a = T !a
  355. foo :: Int -> T Int -> Int
  356. foo 0 t = 0
  357. foo x t | even x = case t of { T n -> foo (x-n) t }
  358. | otherwise = foo (x-1) t
  359. SpecConstr does no specialisation, because the second recursive call
  360. looks like a boxed use of the argument. A pity.
  361. $wfoo_sFw :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
  362. $wfoo_sFw =
  363. \ (ww_sFo [Just L] :: GHC.Prim.Int#) (w_sFq [Just L] :: T.T GHC.Base.Int) ->
  364. case ww_sFo of ds_Xw6 [Just L] {
  365. __DEFAULT ->
  366. case GHC.Prim.remInt# ds_Xw6 2 of wild1_aEF [Dead Just A] {
  367. __DEFAULT -> $wfoo_sFw (GHC.Prim.-# ds_Xw6 1) w_sFq;
  368. 0 ->
  369. case w_sFq of wild_Xy [Just L] { T.T n_ad5 [Just U(L)] ->
  370. case n_ad5 of wild1_aET [Just A] { GHC.Base.I# y_aES [Just L] ->
  371. $wfoo_sFw (GHC.Prim.-# ds_Xw6 y_aES) wild_Xy
  372. } } };
  373. 0 -> 0
  374. Example 2
  375. ~~~~~~~~~
  376. data a :*: b = !a :*: !b
  377. data T a = T !a
  378. foo :: (Int :*: T Int) -> Int
  379. foo (0 :*: t) = 0
  380. foo (x :*: t) | even x = case t of { T n -> foo ((x-n) :*: t) }
  381. | otherwise = foo ((x-1) :*: t)
  382. Very similar to the previous one, except that the parameters are now in
  383. a strict tuple. Before SpecConstr, we have
  384. $wfoo_sG3 :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
  385. $wfoo_sG3 =
  386. \ (ww_sFU [Just L] :: GHC.Prim.Int#) (ww_sFW [Just L] :: T.T
  387. GHC.Base.Int) ->
  388. case ww_sFU of ds_Xws [Just L] {
  389. __DEFAULT ->
  390. case GHC.Prim.remInt# ds_Xws 2 of wild1_aEZ [Dead Just A] {
  391. __DEFAULT ->
  392. case ww_sFW of tpl_B2 [Just L] { T.T a_sFo [Just A] ->
  393. $wfoo_sG3 (GHC.Prim.-# ds_Xws 1) tpl_B2 -- $wfoo1
  394. };
  395. 0 ->
  396. case ww_sFW of wild_XB [Just A] { T.T n_ad7 [Just S(L)] ->
  397. case n_ad7 of wild1_aFd [Just L] { GHC.Base.I# y_aFc [Just L] ->
  398. $wfoo_sG3 (GHC.Prim.-# ds_Xws y_aFc) wild_XB -- $wfoo2
  399. } } };
  400. 0 -> 0 }
  401. We get two specialisations:
  402. "SC:$wfoo1" [0] __forall {a_sFB :: GHC.Base.Int sc_sGC :: GHC.Prim.Int#}
  403. Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int a_sFB)
  404. = Foo.$s$wfoo1 a_sFB sc_sGC ;
  405. "SC:$wfoo2" [0] __forall {y_aFp :: GHC.Prim.Int# sc_sGC :: GHC.Prim.Int#}
  406. Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int (GHC.Base.I# y_aFp))
  407. = Foo.$s$wfoo y_aFp sc_sGC ;
  408. But perhaps the first one isn't good. After all, we know that tpl_B2 is
  409. a T (I# x) really, because T is strict and Int has one constructor. (We can't
  410. unbox the strict fields, becuase T is polymorphic!)
  411. %************************************************************************
  412. %* *
  413. \subsection{Top level wrapper stuff}
  414. %* *
  415. %************************************************************************
  416. \begin{code}
  417. specConstrProgram :: ModGuts -> CoreM ModGuts
  418. specConstrProgram guts
  419. = do
  420. dflags <- getDynFlags
  421. us <- getUniqueSupplyM
  422. annos <- getFirstAnnotations deserializeWithData guts
  423. let binds' = fst $ initUs us (go (initScEnv dflags annos) (mg_binds guts))
  424. return (guts { mg_binds = binds' })
  425. where
  426. go _ [] = return []
  427. go env (bind:binds) = do (env', bind') <- scTopBind env bind
  428. binds' <- go env' binds
  429. return (bind' : binds')
  430. \end{code}
  431. %************************************************************************
  432. %* *
  433. \subsection{Environment: goes downwards}
  434. %* *
  435. %************************************************************************
  436. \begin{code}
  437. data ScEnv = SCE { sc_size :: Maybe Int, -- Size threshold
  438. sc_count :: Maybe Int, -- Max # of specialisations for any one fn
  439. -- See Note [Avoiding exponential blowup]
  440. sc_subst :: Subst, -- Current substitution
  441. -- Maps InIds to OutExprs
  442. sc_how_bound :: HowBoundEnv,
  443. -- Binds interesting non-top-level variables
  444. -- Domain is OutVars (*after* applying the substitution)
  445. sc_vals :: ValueEnv,
  446. -- Domain is OutIds (*after* applying the substitution)
  447. -- Used even for top-level bindings (but not imported ones)
  448. sc_annotations :: UniqFM SpecConstrAnnotation
  449. }
  450. ---------------------
  451. -- As we go, we apply a substitution (sc_subst) to the current term
  452. type InExpr = CoreExpr -- _Before_ applying the subst
  453. type InVar = Var
  454. type OutExpr = CoreExpr -- _After_ applying the subst
  455. type OutId = Id
  456. type OutVar = Var
  457. ---------------------
  458. type HowBoundEnv = VarEnv HowBound -- Domain is OutVars
  459. ---------------------
  460. type ValueEnv = IdEnv Value -- Domain is OutIds
  461. data Value = ConVal AltCon [CoreArg] -- _Saturated_ constructors
  462. -- The AltCon is never DEFAULT
  463. | LambdaVal -- Inlinable lambdas or PAPs
  464. instance Outputable Value where
  465. ppr (ConVal con args) = ppr con <+> interpp'SP args
  466. ppr LambdaVal = ptext (sLit "<Lambda>")
  467. ---------------------
  468. initScEnv :: DynFlags -> UniqFM SpecConstrAnnotation -> ScEnv
  469. initScEnv dflags anns
  470. = SCE { sc_size = specConstrThreshold dflags,
  471. sc_count = specConstrCount dflags,
  472. sc_subst = emptySubst,
  473. sc_how_bound = emptyVarEnv,
  474. sc_vals = emptyVarEnv,
  475. sc_annotations = anns }
  476. data HowBound = RecFun -- These are the recursive functions for which
  477. -- we seek interesting call patterns
  478. | RecArg -- These are those functions' arguments, or their sub-components;
  479. -- we gather occurrence information for these
  480. instance Outputable HowBound where
  481. ppr RecFun = text "RecFun"
  482. ppr RecArg = text "RecArg"
  483. lookupHowBound :: ScEnv -> Id -> Maybe HowBound
  484. lookupHowBound env id = lookupVarEnv (sc_how_bound env) id
  485. scSubstId :: ScEnv -> Id -> CoreExpr
  486. scSubstId env v = lookupIdSubst (text "scSubstId") (sc_subst env) v
  487. scSubstTy :: ScEnv -> Type -> Type
  488. scSubstTy env ty = substTy (sc_subst env) ty
  489. zapScSubst :: ScEnv -> ScEnv
  490. zapScSubst env = env { sc_subst = zapSubstEnv (sc_subst env) }
  491. extendScInScope :: ScEnv -> [Var] -> ScEnv
  492. -- Bring the quantified variables into scope
  493. extendScInScope env qvars = env { sc_subst = extendInScopeList (sc_subst env) qvars }
  494. -- Extend the substitution
  495. extendScSubst :: ScEnv -> Var -> OutExpr -> ScEnv
  496. extendScSubst env var expr = env { sc_subst = extendSubst (sc_subst env) var expr }
  497. extendScSubstList :: ScEnv -> [(Var,OutExpr)] -> ScEnv
  498. extendScSubstList env prs = env { sc_subst = extendSubstList (sc_subst env) prs }
  499. extendHowBound :: ScEnv -> [Var] -> HowBound -> ScEnv
  500. extendHowBound env bndrs how_bound
  501. = env { sc_how_bound = extendVarEnvList (sc_how_bound env)
  502. [(bndr,how_bound) | bndr <- bndrs] }
  503. extendBndrsWith :: HowBound -> ScEnv -> [Var] -> (ScEnv, [Var])
  504. extendBndrsWith how_bound env bndrs
  505. = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndrs')
  506. where
  507. (subst', bndrs') = substBndrs (sc_subst env) bndrs
  508. hb_env' = sc_how_bound env `extendVarEnvList`
  509. [(bndr,how_bound) | bndr <- bndrs']
  510. extendBndrWith :: HowBound -> ScEnv -> Var -> (ScEnv, Var)
  511. extendBndrWith how_bound env bndr
  512. = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndr')
  513. where
  514. (subst', bndr') = substBndr (sc_subst env) bndr
  515. hb_env' = extendVarEnv (sc_how_bound env) bndr' how_bound
  516. extendRecBndrs :: ScEnv -> [Var] -> (ScEnv, [Var])
  517. extendRecBndrs env bndrs = (env { sc_subst = subst' }, bndrs')
  518. where
  519. (subst', bndrs') = substRecBndrs (sc_subst env) bndrs
  520. extendBndr :: ScEnv -> Var -> (ScEnv, Var)
  521. extendBndr env bndr = (env { sc_subst = subst' }, bndr')
  522. where
  523. (subst', bndr') = substBndr (sc_subst env) bndr
  524. extendValEnv :: ScEnv -> Id -> Maybe Value -> ScEnv
  525. extendValEnv env _ Nothing = env
  526. extendValEnv env id (Just cv) = env { sc_vals = extendVarEnv (sc_vals env) id cv }
  527. extendCaseBndrs :: ScEnv -> OutExpr -> OutId -> AltCon -> [Var] -> (ScEnv, [Var])
  528. -- When we encounter
  529. -- case scrut of b
  530. -- C x y -> ...
  531. -- we want to bind b, to (C x y)
  532. -- NB1: Extends only the sc_vals part of the envt
  533. -- NB2: Kill the dead-ness info on the pattern binders x,y, since
  534. -- they are potentially made alive by the [b -> C x y] binding
  535. extendCaseBndrs env scrut case_bndr con alt_bndrs
  536. = (env2, alt_bndrs')
  537. where
  538. live_case_bndr = not (isDeadBinder case_bndr)
  539. env1 | Var v <- scrut = extendValEnv env v cval
  540. | otherwise = env -- See Note [Add scrutinee to ValueEnv too]
  541. env2 | live_case_bndr = extendValEnv env1 case_bndr cval
  542. | otherwise = env1
  543. alt_bndrs' | case scrut of { Var {} -> True; _ -> live_case_bndr }
  544. = map zap alt_bndrs
  545. | otherwise
  546. = alt_bndrs
  547. cval = case con of
  548. DEFAULT -> Nothing
  549. LitAlt {} -> Just (ConVal con [])
  550. DataAlt {} -> Just (ConVal con vanilla_args)
  551. where
  552. vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
  553. varsToCoreExprs alt_bndrs
  554. zap v | isTyCoVar v = v -- See NB2 above
  555. | otherwise = zapIdOccInfo v
  556. decreaseSpecCount :: ScEnv -> Int -> ScEnv
  557. -- See Note [Avoiding exponential blowup]
  558. decreaseSpecCount env n_specs
  559. = env { sc_count = case sc_count env of
  560. Nothing -> Nothing
  561. Just n -> Just (n `div` (n_specs + 1)) }
  562. -- The "+1" takes account of the original function;
  563. -- See Note [Avoiding exponential blowup]
  564. ---------------------------------------------------
  565. -- See Note [SpecConstrAnnotation]
  566. ignoreType :: ScEnv -> Type -> Bool
  567. ignoreAltCon :: ScEnv -> AltCon -> Bool
  568. forceSpecBndr :: ScEnv -> Var -> Bool
  569. #ifndef GHCI
  570. ignoreType _ _ = False
  571. ignoreAltCon _ _ = False
  572. forceSpecBndr _ _ = False
  573. #else /* GHCI */
  574. ignoreAltCon env (DataAlt dc) = ignoreTyCon env (dataConTyCon dc)
  575. ignoreAltCon env (LitAlt lit) = ignoreType env (literalType lit)
  576. ignoreAltCon _ DEFAULT = panic "ignoreAltCon" -- DEFAULT cannot be in a ConVal
  577. ignoreType env ty
  578. = case splitTyConApp_maybe ty of
  579. Just (tycon, _) -> ignoreTyCon env tycon
  580. _ -> False
  581. ignoreTyCon :: ScEnv -> TyCon -> Bool
  582. ignoreTyCon env tycon
  583. = lookupUFM (sc_annotations env) tycon == Just NoSpecConstr
  584. forceSpecBndr env var = forceSpecFunTy env . snd . splitForAllTys . varType $ var
  585. forceSpecFunTy :: ScEnv -> Type -> Bool
  586. forceSpecFunTy env = any (forceSpecArgTy env) . fst . splitFunTys
  587. forceSpecArgTy :: ScEnv -> Type -> Bool
  588. forceSpecArgTy env ty
  589. | Just ty' <- coreView ty = forceSpecArgTy env ty'
  590. forceSpecArgTy env ty
  591. | Just (tycon, tys) <- splitTyConApp_maybe ty
  592. , tycon /= funTyCon
  593. = lookupUFM (sc_annotations env) tycon == Just ForceSpecConstr
  594. || any (forceSpecArgTy env) tys
  595. forceSpecArgTy _ _ = False
  596. #endif /* GHCI */
  597. \end{code}
  598. Note [Add scrutinee to ValueEnv too]
  599. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  600. Consider this:
  601. case x of y
  602. (a,b) -> case b of c
  603. I# v -> ...(f y)...
  604. By the time we get to the call (f y), the ValueEnv
  605. will have a binding for y, and for c
  606. y -> (a,b)
  607. c -> I# v
  608. BUT that's not enough! Looking at the call (f y) we
  609. see that y is pair (a,b), but we also need to know what 'b' is.
  610. So in extendCaseBndrs we must *also* add the binding
  611. b -> I# v
  612. else we lose a useful specialisation for f. This is necessary even
  613. though the simplifier has systematically replaced uses of 'x' with 'y'
  614. and 'b' with 'c' in the code. The use of 'b' in the ValueEnv came
  615. from outside the case. See Trac #4908 for the live example.
  616. Note [Avoiding exponential blowup]
  617. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  618. The sc_count field of the ScEnv says how many times we are prepared to
  619. duplicate a single function. But we must take care with recursive
  620. specialiations. Consider
  621. let $j1 = let $j2 = let $j3 = ...
  622. in
  623. ...$j3...
  624. in
  625. ...$j2...
  626. in
  627. ...$j1...
  628. If we specialise $j1 then in each specialisation (as well as the original)
  629. we can specialise $j2, and similarly $j3. Even if we make just *one*
  630. specialisation of each, becuase we also have the original we'll get 2^n
  631. copies of $j3, which is not good.
  632. So when recursively specialising we divide the sc_count by the number of
  633. copies we are making at this level, including the original.
  634. %************************************************************************
  635. %* *
  636. \subsection{Usage information: flows upwards}
  637. %* *
  638. %************************************************************************
  639. \begin{code}
  640. data ScUsage
  641. = SCU {
  642. scu_calls :: CallEnv, -- Calls
  643. -- The functions are a subset of the
  644. -- RecFuns in the ScEnv
  645. scu_occs :: !(IdEnv ArgOcc) -- Information on argument occurrences
  646. } -- The domain is OutIds
  647. type CallEnv = IdEnv [Call]
  648. type Call = (ValueEnv, [CoreArg])
  649. -- The arguments of the call, together with the
  650. -- env giving the constructor bindings at the call site
  651. nullUsage :: ScUsage
  652. nullUsage = SCU { scu_calls = emptyVarEnv, scu_occs = emptyVarEnv }
  653. combineCalls :: CallEnv -> CallEnv -> CallEnv
  654. combineCalls = plusVarEnv_C (++)
  655. combineUsage :: ScUsage -> ScUsage -> ScUsage
  656. combineUsage u1 u2 = SCU { scu_calls = combineCalls (scu_calls u1) (scu_calls u2),
  657. scu_occs = plusVarEnv_C combineOcc (scu_occs u1) (scu_occs u2) }
  658. combineUsages :: [ScUsage] -> ScUsage
  659. combineUsages [] = nullUsage
  660. combineUsages us = foldr1 combineUsage us
  661. lookupOcc :: ScUsage -> OutVar -> (ScUsage, ArgOcc)
  662. lookupOcc (SCU { scu_calls = sc_calls, scu_occs = sc_occs }) bndr
  663. = (SCU {scu_calls = sc_calls, scu_occs = delVarEnv sc_occs bndr},
  664. lookupVarEnv sc_occs bndr `orElse` NoOcc)
  665. lookupOccs :: ScUsage -> [OutVar] -> (ScUsage, [ArgOcc])
  666. lookupOccs (SCU { scu_calls = sc_calls, scu_occs = sc_occs }) bndrs
  667. = (SCU {scu_calls = sc_calls, scu_occs = delVarEnvList sc_occs bndrs},
  668. [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])
  669. data ArgOcc = NoOcc -- Doesn't occur at all; or a type argument
  670. | UnkOcc -- Used in some unknown way
  671. | ScrutOcc (UniqFM [ArgOcc]) -- See Note [ScrutOcc]
  672. | BothOcc -- Definitely taken apart, *and* perhaps used in some other way
  673. {- Note [ScrutOcc]
  674. An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
  675. is *only* taken apart or applied.
  676. Functions, literal: ScrutOcc emptyUFM
  677. Data constructors: ScrutOcc subs,
  678. where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
  679. The domain of the UniqFM is the Unique of the data constructor
  680. The [ArgOcc] is the occurrences of the *pattern-bound* components
  681. of the data structure. E.g.
  682. data T a = forall b. MkT a b (b->a)
  683. A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
  684. -}
  685. instance Outputable ArgOcc where
  686. ppr (ScrutOcc xs) = ptext (sLit "scrut-occ") <> ppr xs
  687. ppr UnkOcc = ptext (sLit "unk-occ")
  688. ppr BothOcc = ptext (sLit "both-occ")
  689. ppr NoOcc = ptext (sLit "no-occ")
  690. -- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
  691. -- that if the thing is scrutinised anywhere then we get to see that
  692. -- in the overall result, even if it's also used in a boxed way
  693. -- This might be too agressive; see Note [Reboxing] Alternative 3
  694. combineOcc :: ArgOcc -> ArgOcc -> ArgOcc
  695. combineOcc NoOcc occ = occ
  696. combineOcc occ NoOcc = occ
  697. combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
  698. combineOcc _occ (ScrutOcc ys) = ScrutOcc ys
  699. combineOcc (ScrutOcc xs) _occ = ScrutOcc xs
  700. combineOcc UnkOcc UnkOcc = UnkOcc
  701. combineOcc _ _ = BothOcc
  702. combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
  703. combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys
  704. setScrutOcc :: ScEnv -> ScUsage -> OutExpr -> ArgOcc -> ScUsage
  705. -- _Overwrite_ the occurrence info for the scrutinee, if the scrutinee
  706. -- is a variable, and an interesting variable
  707. setScrutOcc env usg (Cast e _) occ = setScrutOcc env usg e occ
  708. setScrutOcc env usg (Note _ e) occ = setScrutOcc env usg e occ
  709. setScrutOcc env usg (Var v) occ
  710. | Just RecArg <- lookupHowBound env v = usg { scu_occs = extendVarEnv (scu_occs usg) v occ }
  711. | otherwise = usg
  712. setScrutOcc _env usg _other _occ -- Catch-all
  713. = usg
  714. conArgOccs :: ArgOcc -> AltCon -> [ArgOcc]
  715. -- Find usage of components of data con; returns [UnkOcc...] if unknown
  716. -- See Note [ScrutOcc] for the extra UnkOccs in the vanilla datacon case
  717. conArgOccs (ScrutOcc fm) (DataAlt dc)
  718. | Just pat_arg_occs <- lookupUFM fm dc
  719. = [UnkOcc | _ <- dataConUnivTyVars dc] ++ pat_arg_occs
  720. conArgOccs _other _con = repeat UnkOcc
  721. \end{code}
  722. %************************************************************************
  723. %* *
  724. \subsection{The main recursive function}
  725. %* *
  726. %************************************************************************
  727. The main recursive function gathers up usage information, and
  728. creates specialised versions of functions.
  729. \begin{code}
  730. scExpr, scExpr' :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
  731. -- The unique supply is needed when we invent
  732. -- a new name for the specialised function and its args
  733. scExpr env e = scExpr' env e
  734. scExpr' env (Var v) = case scSubstId env v of
  735. Var v' -> return (varUsage env v' UnkOcc, Var v')
  736. e' -> scExpr (zapScSubst env) e'
  737. scExpr' env (Type t) = return (nullUsage, Type (scSubstTy env t))
  738. scExpr' _ e@(Lit {}) = return (nullUsage, e)
  739. scExpr' env (Note n e) = do (usg,e') <- scExpr env e
  740. return (usg, Note n e')
  741. scExpr' env (Cast e co) = do (usg, e') <- scExpr env e
  742. return (usg, Cast e' (scSubstTy env co))
  743. scExpr' env e@(App _ _) = scApp env (collectArgs e)
  744. scExpr' env (Lam b e) = do let (env', b') = extendBndr env b
  745. (usg, e') <- scExpr env' e
  746. return (usg, Lam b' e')
  747. scExpr' env (Case scrut b ty alts)
  748. = do { (scrut_usg, scrut') <- scExpr env scrut
  749. ; case isValue (sc_vals env) scrut' of
  750. Just (ConVal con args) -> sc_con_app con args scrut'
  751. _other -> sc_vanilla scrut_usg scrut'
  752. }
  753. where
  754. sc_con_app con args scrut' -- Known constructor; simplify
  755. = do { let (_, bs, rhs) = findAlt con alts
  756. `orElse` (DEFAULT, [], mkImpossibleExpr (coreAltsType alts))
  757. alt_env' = extendScSubstList env ((b,scrut') : bs `zip` trimConArgs con args)
  758. ; scExpr alt_env' rhs }
  759. sc_vanilla scrut_usg scrut' -- Normal case
  760. = do { let (alt_env,b') = extendBndrWith RecArg env b
  761. -- Record RecArg for the components
  762. ; (alt_usgs, alt_occs, alts')
  763. <- mapAndUnzip3M (sc_alt alt_env scrut' b') alts
  764. ; let (alt_usg, b_occ) = lookupOcc (combineUsages alt_usgs) b'
  765. scrut_occ = foldr combineOcc b_occ alt_occs
  766. scrut_usg' = setScrutOcc env scrut_usg scrut' scrut_occ
  767. -- The combined usage of the scrutinee is given
  768. -- by scrut_occ, which is passed to scScrut, which
  769. -- in turn treats a bare-variable scrutinee specially
  770. ; return (alt_usg `combineUsage` scrut_usg',
  771. Case scrut' b' (scSubstTy env ty) alts') }
  772. sc_alt env scrut' b' (con,bs,rhs)
  773. = do { let (env1, bs1) = extendBndrsWith RecArg env bs
  774. (env2, bs2) = extendCaseBndrs env1 scrut' b' con bs1
  775. ; (usg,rhs') <- scExpr env2 rhs
  776. ; let (usg', arg_occs) = lookupOccs usg bs2
  777. scrut_occ = case con of
  778. DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
  779. _ -> ScrutOcc emptyUFM
  780. ; return (usg', scrut_occ, (con, bs2, rhs')) }
  781. scExpr' env (Let (NonRec bndr rhs) body)
  782. | isTyCoVar bndr -- Type-lets may be created by doBeta
  783. = scExpr' (extendScSubst env bndr rhs) body
  784. | otherwise
  785. = do { let (body_env, bndr') = extendBndr env bndr
  786. ; (rhs_usg, rhs_info) <- scRecRhs env (bndr',rhs)
  787. ; let body_env2 = extendHowBound body_env [bndr'] RecFun
  788. -- Note [Local let bindings]
  789. RI _ rhs' _ _ _ = rhs_info
  790. body_env3 = extendValEnv body_env2 bndr' (isValue (sc_vals env) rhs')
  791. ; (body_usg, body') <- scExpr body_env3 body
  792. -- NB: We don't use the ForceSpecConstr mechanism (see
  793. -- Note [Forcing specialisation]) for non-recursive bindings
  794. -- at the moment. I'm not sure if this is the right thing to do.
  795. ; let force_spec = False
  796. ; (spec_usg, specs) <- specialise env force_spec
  797. (scu_calls body_usg)
  798. rhs_info
  799. (SI [] 0 (Just rhs_usg))
  800. ; return (body_usg { scu_calls = scu_calls body_usg `delVarEnv` bndr' }
  801. `combineUsage` spec_usg,
  802. mkLets [NonRec b r | (b,r) <- specInfoBinds rhs_info specs] body')
  803. }
  804. -- A *local* recursive group: see Note [Local recursive groups]
  805. scExpr' env (Let (Rec prs) body)
  806. = do { let (bndrs,rhss) = unzip prs
  807. (rhs_env1,bndrs') = extendRecBndrs env bndrs
  808. rhs_env2 = extendHowBound rhs_env1 bndrs' RecFun
  809. force_spec = any (forceSpecBndr env) bndrs'
  810. -- Note [Forcing specialisation]
  811. ; (rhs_usgs, rhs_infos) <- mapAndUnzipM (scRecRhs rhs_env2) (bndrs' `zip` rhss)
  812. ; (body_usg, body') <- scExpr rhs_env2 body
  813. -- NB: start specLoop from body_usg
  814. ; (spec_usg, specs) <- specLoop rhs_env2 force_spec
  815. (scu_calls body_usg) rhs_infos nullUsage
  816. [SI [] 0 (Just usg) | usg <- rhs_usgs]
  817. -- Do not unconditionally use rhs_usgs.
  818. -- Instead use them only if we find an unspecialised call
  819. -- See Note [Local recursive groups]
  820. ; let all_usg = spec_usg `combineUsage` body_usg
  821. bind' = Rec (concat (zipWith specInfoBinds rhs_infos specs))
  822. ; return (all_usg { scu_calls = scu_calls all_usg `delVarEnvList` bndrs' },
  823. Let bind' body') }
  824. \end{code}
  825. Note [Local let bindings]
  826. ~~~~~~~~~~~~~~~~~~~~~~~~~
  827. It is not uncommon to find this
  828. let $j = \x. <blah> in ...$j True...$j True...
  829. Here $j is an arbitrary let-bound function, but it often comes up for
  830. join points. We might like to specialise $j for its call patterns.
  831. Notice the difference from a letrec, where we look for call patterns
  832. in the *RHS* of the function. Here we look for call patterns in the
  833. *body* of the let.
  834. At one point I predicated this on the RHS mentioning the outer
  835. recursive function, but that's not essential and might even be
  836. harmful. I'm not sure.
  837. \begin{code}
  838. scApp :: ScEnv -> (InExpr, [InExpr]) -> UniqSM (ScUsage, CoreExpr)
  839. scApp env (Var fn, args) -- Function is a variable
  840. = ASSERT( not (null args) )
  841. do { args_w_usgs <- mapM (scExpr env) args
  842. ; let (arg_usgs, args') = unzip args_w_usgs
  843. arg_usg = combineUsages arg_usgs
  844. ; case scSubstId env fn of
  845. fn'@(Lam {}) -> scExpr (zapScSubst env) (doBeta fn' args')
  846. -- Do beta-reduction and try again
  847. Var fn' -> return (arg_usg `combineUsage` fn_usg, mkApps (Var fn') args')
  848. where
  849. fn_usg = case lookupHowBound env fn' of
  850. Just RecFun -> SCU { scu_calls = unitVarEnv fn' [(sc_vals env, args')],
  851. scu_occs = emptyVarEnv }
  852. Just RecArg -> SCU { scu_calls = emptyVarEnv,
  853. scu_occs = unitVarEnv fn' (ScrutOcc emptyUFM) }
  854. Nothing -> nullUsage
  855. other_fn' -> return (arg_usg, mkApps other_fn' args') }
  856. -- NB: doing this ignores any usage info from the substituted
  857. -- function, but I don't think that matters. If it does
  858. -- we can fix it.
  859. where
  860. doBeta :: OutExpr -> [OutExpr] -> OutExpr
  861. -- ToDo: adjust for System IF
  862. doBeta (Lam bndr body) (arg : args) = Let (NonRec bndr arg) (doBeta body args)
  863. doBeta fn args = mkApps fn args
  864. -- The function is almost always a variable, but not always.
  865. -- In particular, if this pass follows float-in,
  866. -- which it may, we can get
  867. -- (let f = ...f... in f) arg1 arg2
  868. scApp env (other_fn, args)
  869. = do { (fn_usg, fn') <- scExpr env other_fn
  870. ; (arg_usgs, args') <- mapAndUnzipM (scExpr env) args
  871. ; return (combineUsages arg_usgs `combineUsage` fn_usg, mkApps fn' args') }
  872. ----------------------
  873. scTopBind :: ScEnv -> CoreBind -> UniqSM (ScEnv, CoreBind)
  874. scTopBind env (Rec prs)
  875. | Just threshold <- sc_size env
  876. , not force_spec
  877. , not (all (couldBeSmallEnoughToInline threshold) rhss)
  878. -- No specialisation
  879. = do { let (rhs_env,bndrs') = extendRecBndrs env bndrs
  880. ; (_, rhss') <- mapAndUnzipM (scExpr rhs_env) rhss
  881. ; return (rhs_env, Rec (bndrs' `zip` rhss')) }
  882. | otherwise -- Do specialisation
  883. = do { let (rhs_env1,bndrs') = extendRecBndrs env bndrs
  884. rhs_env2 = extendHowBound rhs_env1 bndrs' RecFun
  885. ; (rhs_usgs, rhs_infos) <- mapAndUnzipM (scRecRhs rhs_env2) (bndrs' `zip` rhss)
  886. ; let rhs_usg = combineUsages rhs_usgs
  887. ; (_, specs) <- specLoop rhs_env2 force_spec
  888. (scu_calls rhs_usg) rhs_infos nullUsage
  889. [SI [] 0 Nothing | _ <- bndrs]
  890. ; return (rhs_env1, -- For the body of the letrec, delete the RecFun business
  891. Rec (concat (zipWith specInfoBinds rhs_infos specs))) }
  892. where
  893. (bndrs,rhss) = unzip prs
  894. force_spec = any (forceSpecBndr env) bndrs
  895. -- Note [Forcing specialisation]
  896. scTopBind env (NonRec bndr rhs)
  897. = do { (_, rhs') <- scExpr env rhs
  898. ; let (env1, bndr') = extendBndr env bndr
  899. env2 = extendValEnv env1 bndr' (isValue (sc_vals env) rhs')
  900. ; return (env2, NonRec bndr' rhs') }
  901. ----------------------
  902. scRecRhs :: ScEnv -> (OutId, InExpr) -> UniqSM (ScUsage, RhsInfo)
  903. scRecRhs env (bndr,rhs)
  904. = do { let (arg_bndrs,body) = collectBinders rhs
  905. (body_env, arg_bndrs') = extendBndrsWith RecArg env arg_bndrs
  906. ; (body_usg, body') <- scExpr body_env body
  907. ; let (rhs_usg, arg_occs) = lookupOccs body_usg arg_bndrs'
  908. ; return (rhs_usg, RI bndr (mkLams arg_bndrs' body')
  909. arg_bndrs body arg_occs) }
  910. -- The arg_occs says how the visible,
  911. -- lambda-bound binders of the RHS are used
  912. -- (including the TyVar binders)
  913. -- Two pats are the same if they match both ways
  914. ----------------------
  915. specInfoBinds :: RhsInfo -> SpecInfo -> [(Id,CoreExpr)]
  916. specInfoBinds (RI fn new_rhs _ _ _) (SI specs _ _)
  917. = [(id,rhs) | OS _ _ id rhs <- specs] ++
  918. [(fn `addIdSpecialisations` rules, new_rhs)]
  919. where
  920. rules = [r | OS _ r _ _ <- specs]
  921. ----------------------
  922. varUsage :: ScEnv -> OutVar -> ArgOcc -> ScUsage
  923. varUsage env v use
  924. | Just RecArg <- lookupHowBound env v = SCU { scu_calls = emptyVarEnv
  925. , scu_occs = unitVarEnv v use }
  926. | otherwise = nullUsage
  927. \end{code}
  928. %************************************************************************
  929. %* *
  930. The specialiser itself
  931. %* *
  932. %************************************************************************
  933. \begin{code}
  934. data RhsInfo = RI OutId -- The binder
  935. OutExpr -- The new RHS
  936. [InVar] InExpr -- The *original* RHS (\xs.body)
  937. -- Note [Specialise original body]
  938. [ArgOcc] -- Info on how the xs occur in body
  939. data SpecInfo = SI [OneSpec] -- The specialisations we have generated
  940. Int -- Length of specs; used for numbering them
  941. (Maybe ScUsage) -- Nothing => we have generated specialisations
  942. -- from calls in the *original* RHS
  943. -- Just cs => we haven't, and this is the usage
  944. -- of the original RHS
  945. -- See Note [Local recursive groups]
  946. -- One specialisation: Rule plus definition
  947. data OneSpec = OS CallPat -- Call pattern that generated this specialisation
  948. CoreRule -- Rule connecting original id with the specialisation
  949. OutId OutExpr -- Spec id + its rhs
  950. specLoop :: ScEnv
  951. -> Bool -- force specialisation?
  952. -- Note [Forcing specialisation]
  953. -> CallEnv
  954. -> [RhsInfo]
  955. -> ScUsage -> [SpecInfo] -- One per binder; acccumulating parameter
  956. -> UniqSM (ScUsage, [SpecInfo]) -- ...ditto...
  957. specLoop env force_spec all_calls rhs_infos usg_so_far specs_so_far
  958. = do { specs_w_usg <- zipWithM (specialise env force_spec all_calls) rhs_infos specs_so_far
  959. ; let (new_usg_s, all_specs) = unzip specs_w_usg
  960. new_usg = combineUsages new_usg_s
  961. new_calls = scu_calls new_usg
  962. all_usg = usg_so_far `combineUsage` new_usg
  963. ; if isEmptyVarEnv new_calls then
  964. return (all_usg, all_specs)
  965. else
  966. specLoop env force_spec new_calls rhs_infos all_usg all_specs }
  967. specialise
  968. :: ScEnv
  969. -> Bool -- force specialisation?
  970. -- Note [Forcing specialisation]
  971. -> CallEnv -- Info on calls
  972. -> RhsInfo
  973. -> SpecInfo -- Original RHS plus patterns dealt with
  974. -> UniqSM (ScUsage, SpecInfo) -- New specialised versions and their usage
  975. -- Note: the rhs here is the optimised version of the original rhs
  976. -- So when we make a specialised copy of the RHS, we're starting
  977. -- from an RHS whose nested functions have been optimised already.
  978. specialise env force_spec bind_calls (RI fn _ arg_bndrs body arg_occs)
  979. spec_info@(SI specs spec_count mb_unspec)
  980. | not (isBottomingId fn) -- Note [Do not specialise diverging functions]
  981. , not (isNeverActive (idInlineActivation fn)) -- See Note [Transfer activation]
  982. , notNull arg_bndrs -- Only specialise functions
  983. , Just all_calls <- lookupVarEnv bind_calls fn
  984. = do { (boring_call, pats) <- callsToPats env specs arg_occs all_calls
  985. -- ; pprTrace "specialise" (vcat [ ppr fn <+> text "with" <+> int (length pats) <+> text "good patterns"
  986. -- , text "arg_occs" <+> ppr arg_occs
  987. -- , text "calls" <+> ppr all_calls
  988. -- , text "good pats" <+> ppr pats]) $
  989. -- return ()
  990. -- Bale out if too many specialisations
  991. ; let n_pats = length pats
  992. spec_count' = n_pats + spec_count
  993. ; case sc_count env of
  994. Just max | not force_spec && spec_count' > max
  995. -> pprTrace "SpecConstr" msg $
  996. return (nullUsage, spec_info)
  997. where
  998. msg = vcat [ sep [ ptext (sLit "Function") <+> quotes (ppr fn)
  999. , nest 2 (ptext (sLit "has") <+>
  1000. speakNOf spec_count' (ptext (sLit "call pattern")) <> comma <+>
  1001. ptext (sLit "but the limit is") <+> int max) ]
  1002. , ptext (sLit "Use -fspec-constr-count=n to set the bound")
  1003. , extra ]
  1004. extra | not opt_PprStyle_Debug = ptext (sLit "Use -dppr-debug to see specialisations")
  1005. | otherwise = ptext (sLit "Specialisations:") <+> ppr (pats ++ [p | OS p _ _ _ <- specs])
  1006. _normal_case -> do {
  1007. let spec_env = decreaseSpecCount env n_pats
  1008. ; (spec_usgs, new_specs) <- mapAndUnzipM (spec_one spec_env fn arg_bndrs body)
  1009. (pats `zip` [spec_count..])
  1010. -- See Note [Specialise original body]
  1011. ; let spec_usg = combineUsages spec_usgs
  1012. (new_usg, mb_unspec')
  1013. = case mb_unspec of
  1014. Just rhs_usg | boring_call -> (spec_usg `combineUsage` rhs_usg, Nothing)
  1015. _ -> (spec_usg, mb_unspec)
  1016. ; return (new_usg, SI (new_specs ++ specs) spec_count' mb_unspec') } }
  1017. | otherwise
  1018. = return (nullUsage, spec_info) -- The boring case
  1019. ---------------------
  1020. spec_one :: ScEnv
  1021. -> OutId -- Function
  1022. -> [InVar] -- Lambda-binders of RHS; should match patterns
  1023. -> InExpr -- Body of the original function
  1024. -> (CallPat, Int)
  1025. -> UniqSM (ScUsage, OneSpec) -- Rule and binding
  1026. -- spec_one creates a specialised copy of the function, together
  1027. -- with a rule for using it. I'm very proud of how short this
  1028. -- function is, considering what it does :-).
  1029. {-
  1030. Example
  1031. In-scope: a, x::a
  1032. f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
  1033. [c::*, v::(b,c) are presumably bound by the (...) part]
  1034. ==>
  1035. f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
  1036. (...entire body of f...) [b -> (b,c),
  1037. y -> ((:) (a,(b,c)) (x,v) hw)]
  1038. RULE: forall b::* c::*, -- Note, *not* forall a, x
  1039. v::(b,c),
  1040. hw::[(a,(b,c))] .
  1041. f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
  1042. -}
  1043. spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number)
  1044. = do { spec_uniq <- getUniqueUs
  1045. ; let spec_env = extendScSubstList (extendScInScope env qvars)
  1046. (arg_bndrs `zip` pats)
  1047. fn_name = idName fn
  1048. fn_loc = nameSrcSpan fn_name
  1049. spec_occ = mkSpecOcc (nameOccName fn_name)
  1050. rule_name = mkFastString ("SC:" ++ showSDoc (ppr fn <> int rule_number))
  1051. spec_name = mkInternalName spec_uniq spec_occ fn_loc
  1052. -- ; pprTrace "{spec_one" (ppr (sc_count env) <+> ppr fn <+> ppr pats <+> text "-->" <+> ppr spec_name) $
  1053. -- return ()
  1054. -- Specialise the body
  1055. ; (spec_usg, spec_body) <- scExpr spec_env body
  1056. -- ; pprTrace "done spec_one}" (ppr fn) $
  1057. -- return ()
  1058. -- And build the results
  1059. ; let spec_id = mkLocalId spec_name (mkPiTypes spec_lam_args body_ty)
  1060. `setIdStrictness` spec_str -- See Note [Transfer strictness]
  1061. `setIdArity` count isId spec_lam_args
  1062. spec_str = calcSpecStrictness fn spec_lam_args pats
  1063. (spec_lam_args, spec_call_args) = mkWorkerArgs qvars body_ty
  1064. -- Usual w/w hack to avoid generating
  1065. -- a spec_rhs of unlifted type and no args
  1066. spec_rhs = mkLams spec_lam_args spec_body
  1067. body_ty = exprType spec_body
  1068. rule_rhs = mkVarApps (Var spec_id) spec_call_args
  1069. inline_act = idInlineActivation fn
  1070. rule = mkRule True {- Auto -} True {- Local -}
  1071. rule_name inline_act fn_name qvars pats rule_rhs
  1072. -- See Note [Transfer activation]
  1073. ; return (spec_usg, OS call_pat rule spec_id spec_rhs) }
  1074. calcSpecStrictness :: Id -- The original function
  1075. -> [Var] -> [CoreExpr] -- Call pattern
  1076. -> StrictSig -- Strictness of specialised thing
  1077. -- See Note [Transfer strictness]
  1078. calcSpecStrictness fn qvars pats
  1079. = StrictSig (mkTopDmdType spec_dmds TopRes)
  1080. where
  1081. spec_dmds = [ lookupVarEnv dmd_env qv `orElse` lazyDmd | qv <- qvars, isId qv ]
  1082. StrictSig (DmdType _ dmds _) = idStrictness fn
  1083. dmd_env = go emptyVarEnv dmds pats
  1084. go env ds (Type {} : pats) = go env ds pats
  1085. go env (d:ds) (pat : pats) = go (go_one env d pat) ds pats
  1086. go env _ _ = env
  1087. go_one env d (Var v) = extendVarEnv_C both env v d
  1088. go_one env (Box d) e = go_one env d e
  1089. go_one env (Eval (Prod ds)) e
  1090. | (Var _, args) <- collectArgs e = go env ds args
  1091. go_one env _ _ = env
  1092. \end{code}
  1093. Note [Specialise original body]
  1094. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1095. The RhsInfo for a binding keeps the *original* body of the binding. We
  1096. must specialise that, *not* the result of applying specExpr to the RHS
  1097. (which is also kept in RhsInfo). Otherwise we end up specialising a
  1098. specialised RHS, and that can lead directly to exponential behaviour.
  1099. Note [Transfer activation]
  1100. ~~~~~~~~~~~~~~~~~~~~~~~~~~
  1101. This note is for SpecConstr, but exactly the same thing
  1102. happens in the overloading specialiser; see
  1103. Note [Auto-specialisation and RULES] in Specialise.
  1104. In which phase should the specialise-constructor rules be active?
  1105. Originally I made them always-active, but Manuel found that this
  1106. defeated some clever user-written rules. Then I made them active only
  1107. in Phase 0; after all, current

Large files files are truncated, but you can click here to view the full file