PageRenderTime 70ms CodeModel.GetById 20ms RepoModel.GetById 1ms app.codeStats 0ms

/compiler/specialise/SpecConstr.hs

http://github.com/ghc/ghc
Haskell | 2122 lines | 925 code | 258 blank | 939 comment | 33 complexity | d7bb5f9e4c490d2b4f3a783d957d24ef MD5 | raw file
Possible License(s): MIT, BSD-3-Clause, GPL-3.0
  1. {-
  2. ToDo [Oct 2013]
  3. ~~~~~~~~~~~~~~~
  4. 1. Nuke ForceSpecConstr for good (it is subsumed by GHC.Types.SPEC in ghc-prim)
  5. 2. Nuke NoSpecConstr
  6. (c) The GRASP/AQUA Project, Glasgow University, 1992-1998
  7. \section[SpecConstr]{Specialise over constructors}
  8. -}
  9. {-# LANGUAGE CPP #-}
  10. module SpecConstr(
  11. specConstrProgram
  12. #ifdef GHCI
  13. , SpecConstrAnnotation(..)
  14. #endif
  15. ) where
  16. #include "HsVersions.h"
  17. import CoreSyn
  18. import CoreSubst
  19. import CoreUtils
  20. import CoreUnfold ( couldBeSmallEnoughToInline )
  21. import CoreFVs ( exprsFreeVarsList )
  22. import CoreMonad
  23. import Literal ( litIsLifted )
  24. import HscTypes ( ModGuts(..) )
  25. import WwLib ( mkWorkerArgs )
  26. import DataCon
  27. import Coercion hiding( substCo )
  28. import Rules
  29. import Type hiding ( substTy )
  30. import TyCon ( tyConName )
  31. import Id
  32. import PprCore ( pprParendExpr )
  33. import MkCore ( mkImpossibleExpr )
  34. import Var
  35. import VarEnv
  36. import VarSet
  37. import Name
  38. import BasicTypes
  39. import DynFlags ( DynFlags(..) )
  40. import StaticFlags ( opt_PprStyle_Debug )
  41. import Maybes ( orElse, catMaybes, isJust, isNothing )
  42. import Demand
  43. import GHC.Serialized ( deserializeWithData )
  44. import Util
  45. import Pair
  46. import UniqSupply
  47. import Outputable
  48. import FastString
  49. import UniqFM
  50. import MonadUtils
  51. import Control.Monad ( zipWithM )
  52. import Data.List
  53. import PrelNames ( specTyConName )
  54. import Module
  55. -- See Note [Forcing specialisation]
  56. #ifndef GHCI
  57. type SpecConstrAnnotation = ()
  58. #else
  59. import TyCon ( TyCon )
  60. import GHC.Exts( SpecConstrAnnotation(..) )
  61. #endif
  62. {-
  63. -----------------------------------------------------
  64. Game plan
  65. -----------------------------------------------------
  66. Consider
  67. drop n [] = []
  68. drop 0 xs = []
  69. drop n (x:xs) = drop (n-1) xs
  70. After the first time round, we could pass n unboxed. This happens in
  71. numerical code too. Here's what it looks like in Core:
  72. drop n xs = case xs of
  73. [] -> []
  74. (y:ys) -> case n of
  75. I# n# -> case n# of
  76. 0 -> []
  77. _ -> drop (I# (n# -# 1#)) xs
  78. Notice that the recursive call has an explicit constructor as argument.
  79. Noticing this, we can make a specialised version of drop
  80. RULE: drop (I# n#) xs ==> drop' n# xs
  81. drop' n# xs = let n = I# n# in ...orig RHS...
  82. Now the simplifier will apply the specialisation in the rhs of drop', giving
  83. drop' n# xs = case xs of
  84. [] -> []
  85. (y:ys) -> case n# of
  86. 0 -> []
  87. _ -> drop' (n# -# 1#) xs
  88. Much better!
  89. We'd also like to catch cases where a parameter is carried along unchanged,
  90. but evaluated each time round the loop:
  91. f i n = if i>0 || i>n then i else f (i*2) n
  92. Here f isn't strict in n, but we'd like to avoid evaluating it each iteration.
  93. In Core, by the time we've w/wd (f is strict in i) we get
  94. f i# n = case i# ># 0 of
  95. False -> I# i#
  96. True -> case n of { I# n# ->
  97. case i# ># n# of
  98. False -> I# i#
  99. True -> f (i# *# 2#) n
  100. At the call to f, we see that the argument, n is known to be (I# n#),
  101. and n is evaluated elsewhere in the body of f, so we can play the same
  102. trick as above.
  103. Note [Reboxing]
  104. ~~~~~~~~~~~~~~~
  105. We must be careful not to allocate the same constructor twice. Consider
  106. f p = (...(case p of (a,b) -> e)...p...,
  107. ...let t = (r,s) in ...t...(f t)...)
  108. At the recursive call to f, we can see that t is a pair. But we do NOT want
  109. to make a specialised copy:
  110. f' a b = let p = (a,b) in (..., ...)
  111. because now t is allocated by the caller, then r and s are passed to the
  112. recursive call, which allocates the (r,s) pair again.
  113. This happens if
  114. (a) the argument p is used in other than a case-scrutinisation way.
  115. (b) the argument to the call is not a 'fresh' tuple; you have to
  116. look into its unfolding to see that it's a tuple
  117. Hence the "OR" part of Note [Good arguments] below.
  118. ALTERNATIVE 2: pass both boxed and unboxed versions. This no longer saves
  119. allocation, but does perhaps save evals. In the RULE we'd have
  120. something like
  121. f (I# x#) = f' (I# x#) x#
  122. If at the call site the (I# x) was an unfolding, then we'd have to
  123. rely on CSE to eliminate the duplicate allocation.... This alternative
  124. doesn't look attractive enough to pursue.
  125. ALTERNATIVE 3: ignore the reboxing problem. The trouble is that
  126. the conservative reboxing story prevents many useful functions from being
  127. specialised. Example:
  128. foo :: Maybe Int -> Int -> Int
  129. foo (Just m) 0 = 0
  130. foo x@(Just m) n = foo x (n-m)
  131. Here the use of 'x' will clearly not require boxing in the specialised function.
  132. The strictness analyser has the same problem, in fact. Example:
  133. f p@(a,b) = ...
  134. If we pass just 'a' and 'b' to the worker, it might need to rebox the
  135. pair to create (a,b). A more sophisticated analysis might figure out
  136. precisely the cases in which this could happen, but the strictness
  137. analyser does no such analysis; it just passes 'a' and 'b', and hopes
  138. for the best.
  139. So my current choice is to make SpecConstr similarly aggressive, and
  140. ignore the bad potential of reboxing.
  141. Note [Good arguments]
  142. ~~~~~~~~~~~~~~~~~~~~~
  143. So we look for
  144. * A self-recursive function. Ignore mutual recursion for now,
  145. because it's less common, and the code is simpler for self-recursion.
  146. * EITHER
  147. a) At a recursive call, one or more parameters is an explicit
  148. constructor application
  149. AND
  150. That same parameter is scrutinised by a case somewhere in
  151. the RHS of the function
  152. OR
  153. b) At a recursive call, one or more parameters has an unfolding
  154. that is an explicit constructor application
  155. AND
  156. That same parameter is scrutinised by a case somewhere in
  157. the RHS of the function
  158. AND
  159. Those are the only uses of the parameter (see Note [Reboxing])
  160. What to abstract over
  161. ~~~~~~~~~~~~~~~~~~~~~
  162. There's a bit of a complication with type arguments. If the call
  163. site looks like
  164. f p = ...f ((:) [a] x xs)...
  165. then our specialised function look like
  166. f_spec x xs = let p = (:) [a] x xs in ....as before....
  167. This only makes sense if either
  168. a) the type variable 'a' is in scope at the top of f, or
  169. b) the type variable 'a' is an argument to f (and hence fs)
  170. Actually, (a) may hold for value arguments too, in which case
  171. we may not want to pass them. Supose 'x' is in scope at f's
  172. defn, but xs is not. Then we'd like
  173. f_spec xs = let p = (:) [a] x xs in ....as before....
  174. Similarly (b) may hold too. If x is already an argument at the
  175. call, no need to pass it again.
  176. Finally, if 'a' is not in scope at the call site, we could abstract
  177. it as we do the term variables:
  178. f_spec a x xs = let p = (:) [a] x xs in ...as before...
  179. So the grand plan is:
  180. * abstract the call site to a constructor-only pattern
  181. e.g. C x (D (f p) (g q)) ==> C s1 (D s2 s3)
  182. * Find the free variables of the abstracted pattern
  183. * Pass these variables, less any that are in scope at
  184. the fn defn. But see Note [Shadowing] below.
  185. NOTICE that we only abstract over variables that are not in scope,
  186. so we're in no danger of shadowing variables used in "higher up"
  187. in f_spec's RHS.
  188. Note [Shadowing]
  189. ~~~~~~~~~~~~~~~~
  190. In this pass we gather up usage information that may mention variables
  191. that are bound between the usage site and the definition site; or (more
  192. seriously) may be bound to something different at the definition site.
  193. For example:
  194. f x = letrec g y v = let x = ...
  195. in ...(g (a,b) x)...
  196. Since 'x' is in scope at the call site, we may make a rewrite rule that
  197. looks like
  198. RULE forall a,b. g (a,b) x = ...
  199. But this rule will never match, because it's really a different 'x' at
  200. the call site -- and that difference will be manifest by the time the
  201. simplifier gets to it. [A worry: the simplifier doesn't *guarantee*
  202. no-shadowing, so perhaps it may not be distinct?]
  203. Anyway, the rule isn't actually wrong, it's just not useful. One possibility
  204. is to run deShadowBinds before running SpecConstr, but instead we run the
  205. simplifier. That gives the simplest possible program for SpecConstr to
  206. chew on; and it virtually guarantees no shadowing.
  207. Note [Specialising for constant parameters]
  208. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  209. This one is about specialising on a *constant* (but not necessarily
  210. constructor) argument
  211. foo :: Int -> (Int -> Int) -> Int
  212. foo 0 f = 0
  213. foo m f = foo (f m) (+1)
  214. It produces
  215. lvl_rmV :: GHC.Base.Int -> GHC.Base.Int
  216. lvl_rmV =
  217. \ (ds_dlk :: GHC.Base.Int) ->
  218. case ds_dlk of wild_alH { GHC.Base.I# x_alG ->
  219. GHC.Base.I# (GHC.Prim.+# x_alG 1)
  220. T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
  221. GHC.Prim.Int#
  222. T.$wfoo =
  223. \ (ww_sme :: GHC.Prim.Int#) (w_smg :: GHC.Base.Int -> GHC.Base.Int) ->
  224. case ww_sme of ds_Xlw {
  225. __DEFAULT ->
  226. case w_smg (GHC.Base.I# ds_Xlw) of w1_Xmo { GHC.Base.I# ww1_Xmz ->
  227. T.$wfoo ww1_Xmz lvl_rmV
  228. };
  229. 0 -> 0
  230. }
  231. The recursive call has lvl_rmV as its argument, so we could create a specialised copy
  232. with that argument baked in; that is, not passed at all. Now it can perhaps be inlined.
  233. When is this worth it? Call the constant 'lvl'
  234. - If 'lvl' has an unfolding that is a constructor, see if the corresponding
  235. parameter is scrutinised anywhere in the body.
  236. - If 'lvl' has an unfolding that is a inlinable function, see if the corresponding
  237. parameter is applied (...to enough arguments...?)
  238. Also do this is if the function has RULES?
  239. Also
  240. Note [Specialising for lambda parameters]
  241. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  242. foo :: Int -> (Int -> Int) -> Int
  243. foo 0 f = 0
  244. foo m f = foo (f m) (\n -> n-m)
  245. This is subtly different from the previous one in that we get an
  246. explicit lambda as the argument:
  247. T.$wfoo :: GHC.Prim.Int# -> (GHC.Base.Int -> GHC.Base.Int) ->
  248. GHC.Prim.Int#
  249. T.$wfoo =
  250. \ (ww_sm8 :: GHC.Prim.Int#) (w_sma :: GHC.Base.Int -> GHC.Base.Int) ->
  251. case ww_sm8 of ds_Xlr {
  252. __DEFAULT ->
  253. case w_sma (GHC.Base.I# ds_Xlr) of w1_Xmf { GHC.Base.I# ww1_Xmq ->
  254. T.$wfoo
  255. ww1_Xmq
  256. (\ (n_ad3 :: GHC.Base.Int) ->
  257. case n_ad3 of wild_alB { GHC.Base.I# x_alA ->
  258. GHC.Base.I# (GHC.Prim.-# x_alA ds_Xlr)
  259. })
  260. };
  261. 0 -> 0
  262. }
  263. I wonder if SpecConstr couldn't be extended to handle this? After all,
  264. lambda is a sort of constructor for functions and perhaps it already
  265. has most of the necessary machinery?
  266. Furthermore, there's an immediate win, because you don't need to allocate the lambda
  267. at the call site; and if perchance it's called in the recursive call, then you
  268. may avoid allocating it altogether. Just like for constructors.
  269. Looks cool, but probably rare...but it might be easy to implement.
  270. Note [SpecConstr for casts]
  271. ~~~~~~~~~~~~~~~~~~~~~~~~~~~
  272. Consider
  273. data family T a :: *
  274. data instance T Int = T Int
  275. foo n = ...
  276. where
  277. go (T 0) = 0
  278. go (T n) = go (T (n-1))
  279. The recursive call ends up looking like
  280. go (T (I# ...) `cast` g)
  281. So we want to spot the constructor application inside the cast.
  282. That's why we have the Cast case in argToPat
  283. Note [Local recursive groups]
  284. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  285. For a *local* recursive group, we can see all the calls to the
  286. function, so we seed the specialisation loop from the calls in the
  287. body, not from the calls in the RHS. Consider:
  288. bar m n = foo n (n,n) (n,n) (n,n) (n,n)
  289. where
  290. foo n p q r s
  291. | n == 0 = m
  292. | n > 3000 = case p of { (p1,p2) -> foo (n-1) (p2,p1) q r s }
  293. | n > 2000 = case q of { (q1,q2) -> foo (n-1) p (q2,q1) r s }
  294. | n > 1000 = case r of { (r1,r2) -> foo (n-1) p q (r2,r1) s }
  295. | otherwise = case s of { (s1,s2) -> foo (n-1) p q r (s2,s1) }
  296. If we start with the RHSs of 'foo', we get lots and lots of specialisations,
  297. most of which are not needed. But if we start with the (single) call
  298. in the rhs of 'bar' we get exactly one fully-specialised copy, and all
  299. the recursive calls go to this fully-specialised copy. Indeed, the original
  300. function is later collected as dead code. This is very important in
  301. specialising the loops arising from stream fusion, for example in NDP where
  302. we were getting literally hundreds of (mostly unused) specialisations of
  303. a local function.
  304. In a case like the above we end up never calling the original un-specialised
  305. function. (Although we still leave its code around just in case.)
  306. However, if we find any boring calls in the body, including *unsaturated*
  307. ones, such as
  308. letrec foo x y = ....foo...
  309. in map foo xs
  310. then we will end up calling the un-specialised function, so then we *should*
  311. use the calls in the un-specialised RHS as seeds. We call these
  312. "boring call patterns", and callsToPats reports if it finds any of these.
  313. Note [Seeding top-level recursive groups]
  314. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  315. This seeding is done in the binding for seed_calls in specRec.
  316. 1. If all the bindings in a top-level recursive group are local (not
  317. exported), then all the calls are in the rest of the top-level
  318. bindings. This means we can specialise with those call patterns
  319. ONLY, and NOT with the RHSs of the recursive group (exactly like
  320. Note [Local recursive groups])
  321. 2. But if any of the bindings are exported, the function may be called
  322. with any old arguments, so (for lack of anything better) we specialise
  323. based on
  324. (a) the call patterns in the RHS
  325. (b) the call patterns in the rest of the top-level bindings
  326. NB: before Apr 15 we used (a) only, but Dimitrios had an example
  327. where (b) was crucial, so I added that.
  328. Adding (b) also improved nofib allocation results:
  329. multiplier: 4% better
  330. minimax: 2.8% better
  331. Actually in case (2), instead of using the calls from the RHS, it
  332. would be better to specialise in the importing module. We'd need to
  333. add an INLINEABLE pragma to the function, and then it can be
  334. specialised in the importing scope, just as is done for type classes
  335. in Specialise.specImports. This remains to be done (#10346).
  336. Note [Top-level recursive groups]
  337. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  338. To get the call usage information from "the rest of the top level
  339. bindings" (c.f. Note [Seeding top-level recursive groups]), we work
  340. backwards through the top-level bindings so we see the usage before we
  341. get to the binding of the function. Before we can collect the usage
  342. though, we go through all the bindings and add them to the
  343. environment. This is necessary because usage is only tracked for
  344. functions in the environment. These two passes are called
  345. 'go' and 'goEnv'
  346. in specConstrProgram. (Looks a bit revolting to me.)
  347. Note [Do not specialise diverging functions]
  348. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  349. Specialising a function that just diverges is a waste of code.
  350. Furthermore, it broke GHC (simpl014) thus:
  351. {-# STR Sb #-}
  352. f = \x. case x of (a,b) -> f x
  353. If we specialise f we get
  354. f = \x. case x of (a,b) -> fspec a b
  355. But fspec doesn't have decent strictness info. As it happened,
  356. (f x) :: IO t, so the state hack applied and we eta expanded fspec,
  357. and hence f. But now f's strictness is less than its arity, which
  358. breaks an invariant.
  359. Note [Forcing specialisation]
  360. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  361. With stream fusion and in other similar cases, we want to fully
  362. specialise some (but not necessarily all!) loops regardless of their
  363. size and the number of specialisations.
  364. We allow a library to do this, in one of two ways (one which is
  365. deprecated):
  366. 1) Add a parameter of type GHC.Types.SPEC (from ghc-prim) to the loop body.
  367. 2) (Deprecated) Annotate a type with ForceSpecConstr from GHC.Exts,
  368. and then add *that* type as a parameter to the loop body
  369. The reason #2 is deprecated is because it requires GHCi, which isn't
  370. available for things like a cross compiler using stage1.
  371. Here's a (simplified) example from the `vector` package. You may bring
  372. the special 'force specialization' type into scope by saying:
  373. import GHC.Types (SPEC(..))
  374. or by defining your own type (again, deprecated):
  375. data SPEC = SPEC | SPEC2
  376. {-# ANN type SPEC ForceSpecConstr #-}
  377. (Note this is the exact same definition of GHC.Types.SPEC, just
  378. without the annotation.)
  379. After that, you say:
  380. foldl :: (a -> b -> a) -> a -> Stream b -> a
  381. {-# INLINE foldl #-}
  382. foldl f z (Stream step s _) = foldl_loop SPEC z s
  383. where
  384. foldl_loop !sPEC z s = case step s of
  385. Yield x s' -> foldl_loop sPEC (f z x) s'
  386. Skip -> foldl_loop sPEC z s'
  387. Done -> z
  388. SpecConstr will spot the SPEC parameter and always fully specialise
  389. foldl_loop. Note that
  390. * We have to prevent the SPEC argument from being removed by
  391. w/w which is why (a) SPEC is a sum type, and (b) we have to seq on
  392. the SPEC argument.
  393. * And lastly, the SPEC argument is ultimately eliminated by
  394. SpecConstr itself so there is no runtime overhead.
  395. This is all quite ugly; we ought to come up with a better design.
  396. ForceSpecConstr arguments are spotted in scExpr' and scTopBinds which then set
  397. sc_force to True when calling specLoop. This flag does four things:
  398. * Ignore specConstrThreshold, to specialise functions of arbitrary size
  399. (see scTopBind)
  400. * Ignore specConstrCount, to make arbitrary numbers of specialisations
  401. (see specialise)
  402. * Specialise even for arguments that are not scrutinised in the loop
  403. (see argToPat; Trac #4488)
  404. * Only specialise on recursive types a finite number of times
  405. (see is_too_recursive; Trac #5550; Note [Limit recursive specialisation])
  406. This flag is inherited for nested non-recursive bindings (which are likely to
  407. be join points and hence should be fully specialised) but reset for nested
  408. recursive bindings.
  409. What alternatives did I consider? Annotating the loop itself doesn't
  410. work because (a) it is local and (b) it will be w/w'ed and having
  411. w/w propagating annotations somehow doesn't seem like a good idea. The
  412. types of the loop arguments really seem to be the most persistent
  413. thing.
  414. Annotating the types that make up the loop state doesn't work,
  415. either, because (a) it would prevent us from using types like Either
  416. or tuples here, (b) we don't want to restrict the set of types that
  417. can be used in Stream states and (c) some types are fixed by the user
  418. (e.g., the accumulator here) but we still want to specialise as much
  419. as possible.
  420. Alternatives to ForceSpecConstr
  421. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  422. Instead of giving the loop an extra argument of type SPEC, we
  423. also considered *wrapping* arguments in SPEC, thus
  424. data SPEC a = SPEC a | SPEC2
  425. loop = \arg -> case arg of
  426. SPEC state ->
  427. case state of (x,y) -> ... loop (SPEC (x',y')) ...
  428. S2 -> error ...
  429. The idea is that a SPEC argument says "specialise this argument
  430. regardless of whether the function case-analyses it". But this
  431. doesn't work well:
  432. * SPEC must still be a sum type, else the strictness analyser
  433. eliminates it
  434. * But that means that 'loop' won't be strict in its real payload
  435. This loss of strictness in turn screws up specialisation, because
  436. we may end up with calls like
  437. loop (SPEC (case z of (p,q) -> (q,p)))
  438. Without the SPEC, if 'loop' were strict, the case would move out
  439. and we'd see loop applied to a pair. But if 'loop' isn't strict
  440. this doesn't look like a specialisable call.
  441. Note [Limit recursive specialisation]
  442. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  443. It is possible for ForceSpecConstr to cause an infinite loop of specialisation.
  444. Because there is no limit on the number of specialisations, a recursive call with
  445. a recursive constructor as an argument (for example, list cons) will generate
  446. a specialisation for that constructor. If the resulting specialisation also
  447. contains a recursive call with the constructor, this could proceed indefinitely.
  448. For example, if ForceSpecConstr is on:
  449. loop :: [Int] -> [Int] -> [Int]
  450. loop z [] = z
  451. loop z (x:xs) = loop (x:z) xs
  452. this example will create a specialisation for the pattern
  453. loop (a:b) c = loop' a b c
  454. loop' a b [] = (a:b)
  455. loop' a b (x:xs) = loop (x:(a:b)) xs
  456. and a new pattern is found:
  457. loop (a:(b:c)) d = loop'' a b c d
  458. which can continue indefinitely.
  459. Roman's suggestion to fix this was to stop after a couple of times on recursive types,
  460. but still specialising on non-recursive types as much as possible.
  461. To implement this, we count the number of recursive constructors in each
  462. function argument. If the maximum is greater than the specConstrRecursive limit,
  463. do not specialise on that pattern.
  464. This is only necessary when ForceSpecConstr is on: otherwise the specConstrCount
  465. will force termination anyway.
  466. See Trac #5550.
  467. Note [NoSpecConstr]
  468. ~~~~~~~~~~~~~~~~~~~
  469. The ignoreDataCon stuff allows you to say
  470. {-# ANN type T NoSpecConstr #-}
  471. to mean "don't specialise on arguments of this type". It was added
  472. before we had ForceSpecConstr. Lacking ForceSpecConstr we specialised
  473. regardless of size; and then we needed a way to turn that *off*. Now
  474. that we have ForceSpecConstr, this NoSpecConstr is probably redundant.
  475. (Used only for PArray.)
  476. -----------------------------------------------------
  477. Stuff not yet handled
  478. -----------------------------------------------------
  479. Here are notes arising from Roman's work that I don't want to lose.
  480. Example 1
  481. ~~~~~~~~~
  482. data T a = T !a
  483. foo :: Int -> T Int -> Int
  484. foo 0 t = 0
  485. foo x t | even x = case t of { T n -> foo (x-n) t }
  486. | otherwise = foo (x-1) t
  487. SpecConstr does no specialisation, because the second recursive call
  488. looks like a boxed use of the argument. A pity.
  489. $wfoo_sFw :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
  490. $wfoo_sFw =
  491. \ (ww_sFo [Just L] :: GHC.Prim.Int#) (w_sFq [Just L] :: T.T GHC.Base.Int) ->
  492. case ww_sFo of ds_Xw6 [Just L] {
  493. __DEFAULT ->
  494. case GHC.Prim.remInt# ds_Xw6 2 of wild1_aEF [Dead Just A] {
  495. __DEFAULT -> $wfoo_sFw (GHC.Prim.-# ds_Xw6 1) w_sFq;
  496. 0 ->
  497. case w_sFq of wild_Xy [Just L] { T.T n_ad5 [Just U(L)] ->
  498. case n_ad5 of wild1_aET [Just A] { GHC.Base.I# y_aES [Just L] ->
  499. $wfoo_sFw (GHC.Prim.-# ds_Xw6 y_aES) wild_Xy
  500. } } };
  501. 0 -> 0
  502. Example 2
  503. ~~~~~~~~~
  504. data a :*: b = !a :*: !b
  505. data T a = T !a
  506. foo :: (Int :*: T Int) -> Int
  507. foo (0 :*: t) = 0
  508. foo (x :*: t) | even x = case t of { T n -> foo ((x-n) :*: t) }
  509. | otherwise = foo ((x-1) :*: t)
  510. Very similar to the previous one, except that the parameters are now in
  511. a strict tuple. Before SpecConstr, we have
  512. $wfoo_sG3 :: GHC.Prim.Int# -> T.T GHC.Base.Int -> GHC.Prim.Int#
  513. $wfoo_sG3 =
  514. \ (ww_sFU [Just L] :: GHC.Prim.Int#) (ww_sFW [Just L] :: T.T
  515. GHC.Base.Int) ->
  516. case ww_sFU of ds_Xws [Just L] {
  517. __DEFAULT ->
  518. case GHC.Prim.remInt# ds_Xws 2 of wild1_aEZ [Dead Just A] {
  519. __DEFAULT ->
  520. case ww_sFW of tpl_B2 [Just L] { T.T a_sFo [Just A] ->
  521. $wfoo_sG3 (GHC.Prim.-# ds_Xws 1) tpl_B2 -- $wfoo1
  522. };
  523. 0 ->
  524. case ww_sFW of wild_XB [Just A] { T.T n_ad7 [Just S(L)] ->
  525. case n_ad7 of wild1_aFd [Just L] { GHC.Base.I# y_aFc [Just L] ->
  526. $wfoo_sG3 (GHC.Prim.-# ds_Xws y_aFc) wild_XB -- $wfoo2
  527. } } };
  528. 0 -> 0 }
  529. We get two specialisations:
  530. "SC:$wfoo1" [0] __forall {a_sFB :: GHC.Base.Int sc_sGC :: GHC.Prim.Int#}
  531. Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int a_sFB)
  532. = Foo.$s$wfoo1 a_sFB sc_sGC ;
  533. "SC:$wfoo2" [0] __forall {y_aFp :: GHC.Prim.Int# sc_sGC :: GHC.Prim.Int#}
  534. Foo.$wfoo sc_sGC (Foo.T @ GHC.Base.Int (GHC.Base.I# y_aFp))
  535. = Foo.$s$wfoo y_aFp sc_sGC ;
  536. But perhaps the first one isn't good. After all, we know that tpl_B2 is
  537. a T (I# x) really, because T is strict and Int has one constructor. (We can't
  538. unbox the strict fields, because T is polymorphic!)
  539. ************************************************************************
  540. * *
  541. \subsection{Top level wrapper stuff}
  542. * *
  543. ************************************************************************
  544. -}
  545. specConstrProgram :: ModGuts -> CoreM ModGuts
  546. specConstrProgram guts
  547. = do
  548. dflags <- getDynFlags
  549. us <- getUniqueSupplyM
  550. annos <- getFirstAnnotations deserializeWithData guts
  551. this_mod <- getModule
  552. let binds' = reverse $ fst $ initUs us $ do
  553. -- Note [Top-level recursive groups]
  554. (env, binds) <- goEnv (initScEnv dflags this_mod annos)
  555. (mg_binds guts)
  556. -- binds is identical to (mg_binds guts), except that the
  557. -- binders on the LHS have been replaced by extendBndr
  558. -- (SPJ this seems like overkill; I don't think the binders
  559. -- will change at all; and we don't substitute in the RHSs anyway!!)
  560. go env nullUsage (reverse binds)
  561. return (guts { mg_binds = binds' })
  562. where
  563. -- See Note [Top-level recursive groups]
  564. goEnv env [] = return (env, [])
  565. goEnv env (bind:binds) = do (env', bind') <- scTopBindEnv env bind
  566. (env'', binds') <- goEnv env' binds
  567. return (env'', bind' : binds')
  568. -- Arg list of bindings is in reverse order
  569. go _ _ [] = return []
  570. go env usg (bind:binds) = do (usg', bind') <- scTopBind env usg bind
  571. binds' <- go env usg' binds
  572. return (bind' : binds')
  573. {-
  574. ************************************************************************
  575. * *
  576. \subsection{Environment: goes downwards}
  577. * *
  578. ************************************************************************
  579. Note [Work-free values only in environment]
  580. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  581. The sc_vals field keeps track of in-scope value bindings, so
  582. that if we come across (case x of Just y ->...) we can reduce the
  583. case from knowing that x is bound to a pair.
  584. But only *work-free* values are ok here. For example if the envt had
  585. x -> Just (expensive v)
  586. then we do NOT want to expand to
  587. let y = expensive v in ...
  588. because the x-binding still exists and we've now duplicated (expensive v).
  589. This seldom happens because let-bound constructor applications are
  590. ANF-ised, but it can happen as a result of on-the-fly transformations in
  591. SpecConstr itself. Here is Trac #7865:
  592. let {
  593. a'_shr =
  594. case xs_af8 of _ {
  595. [] -> acc_af6;
  596. : ds_dgt [Dmd=<L,A>] ds_dgu [Dmd=<L,A>] ->
  597. (expensive x_af7, x_af7
  598. } } in
  599. let {
  600. ds_sht =
  601. case a'_shr of _ { (p'_afd, q'_afe) ->
  602. TSpecConstr_DoubleInline.recursive
  603. (GHC.Types.: @ GHC.Types.Int x_af7 wild_X6) (q'_afe, p'_afd)
  604. } } in
  605. When processed knowing that xs_af8 was bound to a cons, we simplify to
  606. a'_shr = (expensive x_af7, x_af7)
  607. and we do NOT want to inline that at the occurrence of a'_shr in ds_sht.
  608. (There are other occurrences of a'_shr.) No no no.
  609. It would be possible to do some on-the-fly ANF-ising, so that a'_shr turned
  610. into a work-free value again, thus
  611. a1 = expensive x_af7
  612. a'_shr = (a1, x_af7)
  613. but that's more work, so until its shown to be important I'm going to
  614. leave it for now.
  615. -}
  616. data ScEnv = SCE { sc_dflags :: DynFlags,
  617. sc_module :: !Module,
  618. sc_size :: Maybe Int, -- Size threshold
  619. sc_count :: Maybe Int, -- Max # of specialisations for any one fn
  620. -- See Note [Avoiding exponential blowup]
  621. sc_recursive :: Int, -- Max # of specialisations over recursive type.
  622. -- Stops ForceSpecConstr from diverging.
  623. sc_force :: Bool, -- Force specialisation?
  624. -- See Note [Forcing specialisation]
  625. sc_subst :: Subst, -- Current substitution
  626. -- Maps InIds to OutExprs
  627. sc_how_bound :: HowBoundEnv,
  628. -- Binds interesting non-top-level variables
  629. -- Domain is OutVars (*after* applying the substitution)
  630. sc_vals :: ValueEnv,
  631. -- Domain is OutIds (*after* applying the substitution)
  632. -- Used even for top-level bindings (but not imported ones)
  633. -- The range of the ValueEnv is *work-free* values
  634. -- such as (\x. blah), or (Just v)
  635. -- but NOT (Just (expensive v))
  636. -- See Note [Work-free values only in environment]
  637. sc_annotations :: UniqFM SpecConstrAnnotation
  638. }
  639. ---------------------
  640. -- As we go, we apply a substitution (sc_subst) to the current term
  641. type InExpr = CoreExpr -- _Before_ applying the subst
  642. type InVar = Var
  643. type OutExpr = CoreExpr -- _After_ applying the subst
  644. type OutId = Id
  645. type OutVar = Var
  646. ---------------------
  647. type HowBoundEnv = VarEnv HowBound -- Domain is OutVars
  648. ---------------------
  649. type ValueEnv = IdEnv Value -- Domain is OutIds
  650. data Value = ConVal AltCon [CoreArg] -- _Saturated_ constructors
  651. -- The AltCon is never DEFAULT
  652. | LambdaVal -- Inlinable lambdas or PAPs
  653. instance Outputable Value where
  654. ppr (ConVal con args) = ppr con <+> interpp'SP args
  655. ppr LambdaVal = text "<Lambda>"
  656. ---------------------
  657. initScEnv :: DynFlags -> Module -> UniqFM SpecConstrAnnotation -> ScEnv
  658. initScEnv dflags this_mod anns
  659. = SCE { sc_dflags = dflags,
  660. sc_module = this_mod,
  661. sc_size = specConstrThreshold dflags,
  662. sc_count = specConstrCount dflags,
  663. sc_recursive = specConstrRecursive dflags,
  664. sc_force = False,
  665. sc_subst = emptySubst,
  666. sc_how_bound = emptyVarEnv,
  667. sc_vals = emptyVarEnv,
  668. sc_annotations = anns }
  669. data HowBound = RecFun -- These are the recursive functions for which
  670. -- we seek interesting call patterns
  671. | RecArg -- These are those functions' arguments, or their sub-components;
  672. -- we gather occurrence information for these
  673. instance Outputable HowBound where
  674. ppr RecFun = text "RecFun"
  675. ppr RecArg = text "RecArg"
  676. scForce :: ScEnv -> Bool -> ScEnv
  677. scForce env b = env { sc_force = b }
  678. lookupHowBound :: ScEnv -> Id -> Maybe HowBound
  679. lookupHowBound env id = lookupVarEnv (sc_how_bound env) id
  680. scSubstId :: ScEnv -> Id -> CoreExpr
  681. scSubstId env v = lookupIdSubst (text "scSubstId") (sc_subst env) v
  682. scSubstTy :: ScEnv -> Type -> Type
  683. scSubstTy env ty = substTy (sc_subst env) ty
  684. scSubstCo :: ScEnv -> Coercion -> Coercion
  685. scSubstCo env co = substCo (sc_subst env) co
  686. zapScSubst :: ScEnv -> ScEnv
  687. zapScSubst env = env { sc_subst = zapSubstEnv (sc_subst env) }
  688. extendScInScope :: ScEnv -> [Var] -> ScEnv
  689. -- Bring the quantified variables into scope
  690. extendScInScope env qvars = env { sc_subst = extendInScopeList (sc_subst env) qvars }
  691. -- Extend the substitution
  692. extendScSubst :: ScEnv -> Var -> OutExpr -> ScEnv
  693. extendScSubst env var expr = env { sc_subst = extendSubst (sc_subst env) var expr }
  694. extendScSubstList :: ScEnv -> [(Var,OutExpr)] -> ScEnv
  695. extendScSubstList env prs = env { sc_subst = extendSubstList (sc_subst env) prs }
  696. extendHowBound :: ScEnv -> [Var] -> HowBound -> ScEnv
  697. extendHowBound env bndrs how_bound
  698. = env { sc_how_bound = extendVarEnvList (sc_how_bound env)
  699. [(bndr,how_bound) | bndr <- bndrs] }
  700. extendBndrsWith :: HowBound -> ScEnv -> [Var] -> (ScEnv, [Var])
  701. extendBndrsWith how_bound env bndrs
  702. = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndrs')
  703. where
  704. (subst', bndrs') = substBndrs (sc_subst env) bndrs
  705. hb_env' = sc_how_bound env `extendVarEnvList`
  706. [(bndr,how_bound) | bndr <- bndrs']
  707. extendBndrWith :: HowBound -> ScEnv -> Var -> (ScEnv, Var)
  708. extendBndrWith how_bound env bndr
  709. = (env { sc_subst = subst', sc_how_bound = hb_env' }, bndr')
  710. where
  711. (subst', bndr') = substBndr (sc_subst env) bndr
  712. hb_env' = extendVarEnv (sc_how_bound env) bndr' how_bound
  713. extendRecBndrs :: ScEnv -> [Var] -> (ScEnv, [Var])
  714. extendRecBndrs env bndrs = (env { sc_subst = subst' }, bndrs')
  715. where
  716. (subst', bndrs') = substRecBndrs (sc_subst env) bndrs
  717. extendBndr :: ScEnv -> Var -> (ScEnv, Var)
  718. extendBndr env bndr = (env { sc_subst = subst' }, bndr')
  719. where
  720. (subst', bndr') = substBndr (sc_subst env) bndr
  721. extendValEnv :: ScEnv -> Id -> Maybe Value -> ScEnv
  722. extendValEnv env _ Nothing = env
  723. extendValEnv env id (Just cv)
  724. | valueIsWorkFree cv -- Don't duplicate work!! Trac #7865
  725. = env { sc_vals = extendVarEnv (sc_vals env) id cv }
  726. extendValEnv env _ _ = env
  727. extendCaseBndrs :: ScEnv -> OutExpr -> OutId -> AltCon -> [Var] -> (ScEnv, [Var])
  728. -- When we encounter
  729. -- case scrut of b
  730. -- C x y -> ...
  731. -- we want to bind b, to (C x y)
  732. -- NB1: Extends only the sc_vals part of the envt
  733. -- NB2: Kill the dead-ness info on the pattern binders x,y, since
  734. -- they are potentially made alive by the [b -> C x y] binding
  735. extendCaseBndrs env scrut case_bndr con alt_bndrs
  736. = (env2, alt_bndrs')
  737. where
  738. live_case_bndr = not (isDeadBinder case_bndr)
  739. env1 | Var v <- stripTicksTopE (const True) scrut
  740. = extendValEnv env v cval
  741. | otherwise = env -- See Note [Add scrutinee to ValueEnv too]
  742. env2 | live_case_bndr = extendValEnv env1 case_bndr cval
  743. | otherwise = env1
  744. alt_bndrs' | case scrut of { Var {} -> True; _ -> live_case_bndr }
  745. = map zap alt_bndrs
  746. | otherwise
  747. = alt_bndrs
  748. cval = case con of
  749. DEFAULT -> Nothing
  750. LitAlt {} -> Just (ConVal con [])
  751. DataAlt {} -> Just (ConVal con vanilla_args)
  752. where
  753. vanilla_args = map Type (tyConAppArgs (idType case_bndr)) ++
  754. varsToCoreExprs alt_bndrs
  755. zap v | isTyVar v = v -- See NB2 above
  756. | otherwise = zapIdOccInfo v
  757. decreaseSpecCount :: ScEnv -> Int -> ScEnv
  758. -- See Note [Avoiding exponential blowup]
  759. decreaseSpecCount env n_specs
  760. = env { sc_count = case sc_count env of
  761. Nothing -> Nothing
  762. Just n -> Just (n `div` (n_specs + 1)) }
  763. -- The "+1" takes account of the original function;
  764. -- See Note [Avoiding exponential blowup]
  765. ---------------------------------------------------
  766. -- See Note [Forcing specialisation]
  767. ignoreType :: ScEnv -> Type -> Bool
  768. ignoreDataCon :: ScEnv -> DataCon -> Bool
  769. forceSpecBndr :: ScEnv -> Var -> Bool
  770. #ifndef GHCI
  771. ignoreType _ _ = False
  772. ignoreDataCon _ _ = False
  773. #else /* GHCI */
  774. ignoreDataCon env dc = ignoreTyCon env (dataConTyCon dc)
  775. ignoreType env ty
  776. = case tyConAppTyCon_maybe ty of
  777. Just tycon -> ignoreTyCon env tycon
  778. _ -> False
  779. ignoreTyCon :: ScEnv -> TyCon -> Bool
  780. ignoreTyCon env tycon
  781. = lookupUFM (sc_annotations env) tycon == Just NoSpecConstr
  782. #endif /* GHCI */
  783. forceSpecBndr env var = forceSpecFunTy env . snd . splitForAllTys . varType $ var
  784. forceSpecFunTy :: ScEnv -> Type -> Bool
  785. forceSpecFunTy env = any (forceSpecArgTy env) . fst . splitFunTys
  786. forceSpecArgTy :: ScEnv -> Type -> Bool
  787. forceSpecArgTy env ty
  788. | Just ty' <- coreView ty = forceSpecArgTy env ty'
  789. forceSpecArgTy env ty
  790. | Just (tycon, tys) <- splitTyConApp_maybe ty
  791. , tycon /= funTyCon
  792. = tyConName tycon == specTyConName
  793. #ifdef GHCI
  794. || lookupUFM (sc_annotations env) tycon == Just ForceSpecConstr
  795. #endif
  796. || any (forceSpecArgTy env) tys
  797. forceSpecArgTy _ _ = False
  798. {-
  799. Note [Add scrutinee to ValueEnv too]
  800. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  801. Consider this:
  802. case x of y
  803. (a,b) -> case b of c
  804. I# v -> ...(f y)...
  805. By the time we get to the call (f y), the ValueEnv
  806. will have a binding for y, and for c
  807. y -> (a,b)
  808. c -> I# v
  809. BUT that's not enough! Looking at the call (f y) we
  810. see that y is pair (a,b), but we also need to know what 'b' is.
  811. So in extendCaseBndrs we must *also* add the binding
  812. b -> I# v
  813. else we lose a useful specialisation for f. This is necessary even
  814. though the simplifier has systematically replaced uses of 'x' with 'y'
  815. and 'b' with 'c' in the code. The use of 'b' in the ValueEnv came
  816. from outside the case. See Trac #4908 for the live example.
  817. Note [Avoiding exponential blowup]
  818. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  819. The sc_count field of the ScEnv says how many times we are prepared to
  820. duplicate a single function. But we must take care with recursive
  821. specialisations. Consider
  822. let $j1 = let $j2 = let $j3 = ...
  823. in
  824. ...$j3...
  825. in
  826. ...$j2...
  827. in
  828. ...$j1...
  829. If we specialise $j1 then in each specialisation (as well as the original)
  830. we can specialise $j2, and similarly $j3. Even if we make just *one*
  831. specialisation of each, because we also have the original we'll get 2^n
  832. copies of $j3, which is not good.
  833. So when recursively specialising we divide the sc_count by the number of
  834. copies we are making at this level, including the original.
  835. ************************************************************************
  836. * *
  837. \subsection{Usage information: flows upwards}
  838. * *
  839. ************************************************************************
  840. -}
  841. data ScUsage
  842. = SCU {
  843. scu_calls :: CallEnv, -- Calls
  844. -- The functions are a subset of the
  845. -- RecFuns in the ScEnv
  846. scu_occs :: !(IdEnv ArgOcc) -- Information on argument occurrences
  847. } -- The domain is OutIds
  848. type CallEnv = IdEnv [Call]
  849. data Call = Call Id [CoreArg] ValueEnv
  850. -- The arguments of the call, together with the
  851. -- env giving the constructor bindings at the call site
  852. -- We keep the function mainly for debug output
  853. instance Outputable ScUsage where
  854. ppr (SCU { scu_calls = calls, scu_occs = occs })
  855. = text "SCU" <+> braces (sep [ ptext (sLit "calls =") <+> ppr calls
  856. , text "occs =" <+> ppr occs ])
  857. instance Outputable Call where
  858. ppr (Call fn args _) = ppr fn <+> fsep (map pprParendExpr args)
  859. nullUsage :: ScUsage
  860. nullUsage = SCU { scu_calls = emptyVarEnv, scu_occs = emptyVarEnv }
  861. combineCalls :: CallEnv -> CallEnv -> CallEnv
  862. combineCalls = plusVarEnv_C (++)
  863. where
  864. -- plus cs ds | length res > 1
  865. -- = pprTrace "combineCalls" (vcat [ text "cs:" <+> ppr cs
  866. -- , text "ds:" <+> ppr ds])
  867. -- res
  868. -- | otherwise = res
  869. -- where
  870. -- res = cs ++ ds
  871. combineUsage :: ScUsage -> ScUsage -> ScUsage
  872. combineUsage u1 u2 = SCU { scu_calls = combineCalls (scu_calls u1) (scu_calls u2),
  873. scu_occs = plusVarEnv_C combineOcc (scu_occs u1) (scu_occs u2) }
  874. combineUsages :: [ScUsage] -> ScUsage
  875. combineUsages [] = nullUsage
  876. combineUsages us = foldr1 combineUsage us
  877. lookupOccs :: ScUsage -> [OutVar] -> (ScUsage, [ArgOcc])
  878. lookupOccs (SCU { scu_calls = sc_calls, scu_occs = sc_occs }) bndrs
  879. = (SCU {scu_calls = sc_calls, scu_occs = delVarEnvList sc_occs bndrs},
  880. [lookupVarEnv sc_occs b `orElse` NoOcc | b <- bndrs])
  881. data ArgOcc = NoOcc -- Doesn't occur at all; or a type argument
  882. | UnkOcc -- Used in some unknown way
  883. | ScrutOcc -- See Note [ScrutOcc]
  884. (DataConEnv [ArgOcc]) -- How the sub-components are used
  885. type DataConEnv a = UniqFM a -- Keyed by DataCon
  886. {- Note [ScrutOcc]
  887. ~~~~~~~~~~~~~~~~~~~
  888. An occurrence of ScrutOcc indicates that the thing, or a `cast` version of the thing,
  889. is *only* taken apart or applied.
  890. Functions, literal: ScrutOcc emptyUFM
  891. Data constructors: ScrutOcc subs,
  892. where (subs :: UniqFM [ArgOcc]) gives usage of the *pattern-bound* components,
  893. The domain of the UniqFM is the Unique of the data constructor
  894. The [ArgOcc] is the occurrences of the *pattern-bound* components
  895. of the data structure. E.g.
  896. data T a = forall b. MkT a b (b->a)
  897. A pattern binds b, x::a, y::b, z::b->a, but not 'a'!
  898. -}
  899. instance Outputable ArgOcc where
  900. ppr (ScrutOcc xs) = text "scrut-occ" <> ppr xs
  901. ppr UnkOcc = text "unk-occ"
  902. ppr NoOcc = text "no-occ"
  903. evalScrutOcc :: ArgOcc
  904. evalScrutOcc = ScrutOcc emptyUFM
  905. -- Experimentally, this vesion of combineOcc makes ScrutOcc "win", so
  906. -- that if the thing is scrutinised anywhere then we get to see that
  907. -- in the overall result, even if it's also used in a boxed way
  908. -- This might be too aggressive; see Note [Reboxing] Alternative 3
  909. combineOcc :: ArgOcc -> ArgOcc -> ArgOcc
  910. combineOcc NoOcc occ = occ
  911. combineOcc occ NoOcc = occ
  912. combineOcc (ScrutOcc xs) (ScrutOcc ys) = ScrutOcc (plusUFM_C combineOccs xs ys)
  913. combineOcc UnkOcc (ScrutOcc ys) = ScrutOcc ys
  914. combineOcc (ScrutOcc xs) UnkOcc = ScrutOcc xs
  915. combineOcc UnkOcc UnkOcc = UnkOcc
  916. combineOccs :: [ArgOcc] -> [ArgOcc] -> [ArgOcc]
  917. combineOccs xs ys = zipWithEqual "combineOccs" combineOcc xs ys
  918. setScrutOcc :: ScEnv -> ScUsage -> OutExpr -> ArgOcc -> ScUsage
  919. -- _Overwrite_ the occurrence info for the scrutinee, if the scrutinee
  920. -- is a variable, and an interesting variable
  921. setScrutOcc env usg (Cast e _) occ = setScrutOcc env usg e occ
  922. setScrutOcc env usg (Tick _ e) occ = setScrutOcc env usg e occ
  923. setScrutOcc env usg (Var v) occ
  924. | Just RecArg <- lookupHowBound env v = usg { scu_occs = extendVarEnv (scu_occs usg) v occ }
  925. | otherwise = usg
  926. setScrutOcc _env usg _other _occ -- Catch-all
  927. = usg
  928. {-
  929. ************************************************************************
  930. * *
  931. \subsection{The main recursive function}
  932. * *
  933. ************************************************************************
  934. The main recursive function gathers up usage information, and
  935. creates specialised versions of functions.
  936. -}
  937. scExpr, scExpr' :: ScEnv -> CoreExpr -> UniqSM (ScUsage, CoreExpr)
  938. -- The unique supply is needed when we invent
  939. -- a new name for the specialised function and its args
  940. scExpr env e = scExpr' env e
  941. scExpr' env (Var v) = case scSubstId env v of
  942. Var v' -> return (mkVarUsage env v' [], Var v')
  943. e' -> scExpr (zapScSubst env) e'
  944. scExpr' env (Type t) = return (nullUsage, Type (scSubstTy env t))
  945. scExpr' env (Coercion c) = return (nullUsage, Coercion (scSubstCo env c))
  946. scExpr' _ e@(Lit {}) = return (nullUsage, e)
  947. scExpr' env (Tick t e) = do (usg, e') <- scExpr env e
  948. return (usg, Tick t e')
  949. scExpr' env (Cast e co) = do (usg, e') <- scExpr env e
  950. return (usg, mkCast e' (scSubstCo env co))
  951. -- Important to use mkCast here
  952. -- See Note [SpecConstr call patterns]
  953. scExpr' env e@(App _ _) = scApp env (collectArgs e)
  954. scExpr' env (Lam b e) = do let (env', b') = extendBndr env b
  955. (usg, e') <- scExpr env' e
  956. return (usg, Lam b' e')
  957. scExpr' env (Case scrut b ty alts)
  958. = do { (scrut_usg, scrut') <- scExpr env scrut
  959. ; case isValue (sc_vals env) scrut' of
  960. Just (ConVal con args) -> sc_con_app con args scrut'
  961. _other -> sc_vanilla scrut_usg scrut'
  962. }
  963. where
  964. sc_con_app con args scrut' -- Known constructor; simplify
  965. = do { let (_, bs, rhs) = findAlt con alts
  966. `orElse` (DEFAULT, [], mkImpossibleExpr ty)
  967. alt_env' = extendScSubstList env ((b,scrut') : bs `zip` trimConArgs con args)
  968. ; scExpr alt_env' rhs }
  969. sc_vanilla scrut_usg scrut' -- Normal case
  970. = do { let (alt_env,b') = extendBndrWith RecArg env b
  971. -- Record RecArg for the components
  972. ; (alt_usgs, alt_occs, alts')
  973. <- mapAndUnzip3M (sc_alt alt_env scrut' b') alts
  974. ; let scrut_occ = foldr combineOcc NoOcc alt_occs
  975. scrut_usg' = setScrutOcc env scrut_usg scrut' scrut_occ
  976. -- The combined usage of the scrutinee is given
  977. -- by scrut_occ, which is passed to scScrut, which
  978. -- in turn treats a bare-variable scrutinee specially
  979. ; return (foldr combineUsage scrut_usg' alt_usgs,
  980. Case scrut' b' (scSubstTy env ty) alts') }
  981. sc_alt env scrut' b' (con,bs,rhs)
  982. = do { let (env1, bs1) = extendBndrsWith RecArg env bs
  983. (env2, bs2) = extendCaseBndrs env1 scrut' b' con bs1
  984. ; (usg, rhs') <- scExpr env2 rhs
  985. ; let (usg', b_occ:arg_occs) = lookupOccs usg (b':bs2)
  986. scrut_occ = case con of
  987. DataAlt dc -> ScrutOcc (unitUFM dc arg_occs)
  988. _ -> ScrutOcc emptyUFM
  989. ; return (usg', b_occ `combineOcc` scrut_occ, (con, bs2, rhs')) }
  990. scExpr' env (Let (NonRec bndr rhs) body)
  991. | isTyVar bndr -- Type-lets may be created by doBeta
  992. = scExpr' (extendScSubst env bndr rhs) body
  993. | otherwise
  994. = do { let (body_env, bndr') = extendBndr env bndr
  995. ; rhs_info <- scRecRhs env (bndr',rhs)
  996. ; let body_env2 = extendHowBound body_env [bndr'] RecFun
  997. -- Note [Local let bindings]
  998. rhs' = ri_new_rhs rhs_info
  999. body_env3 = extendValEnv body_env2 bndr' (isValue (sc_vals env) rhs')
  1000. ; (body_usg, body') <- scExpr body_env3 body
  1001. -- NB: For non-recursive bindings we inherit sc_force flag from
  1002. -- the parent function (see Note [Forcing specialisation])
  1003. ; (spec_usg, specs) <- specNonRec env body_usg rhs_info
  1004. ; return (body_usg { scu_calls = scu_calls body_usg `delVarEnv` bndr' }
  1005. `combineUsage` spec_usg, -- Note [spec_usg includes rhs_usg]
  1006. mkLets [NonRec b r | (b,r) <- ruleInfoBinds rhs_info specs] body')
  1007. }
  1008. -- A *local* recursive group: see Note [Local recursive groups]
  1009. scExpr' env (Let (Rec prs) body)
  1010. = do { let (bndrs,rhss) = unzip prs
  1011. (rhs_env1,bndrs') = extendRecBndrs env bndrs
  1012. rhs_env2 = extendHowBound rhs_env1 bndrs' RecFun
  1013. force_spec = any (forceSpecBndr env) bndrs'
  1014. -- Note [Forcing specialisation]
  1015. ; rhs_infos <- mapM (scRecRhs rhs_env2) (bndrs' `zip` rhss)
  1016. ; (body_usg, body') <- scExpr rhs_env2 body
  1017. -- NB: start specLoop from body_usg
  1018. ; (spec_usg, specs) <- specRec NotTopLevel (scForce rhs_env2 force_spec)
  1019. body_usg rhs_infos
  1020. -- Do not unconditionally generate specialisations from rhs_usgs
  1021. -- Instead use them only if we find an unspecialised call
  1022. -- See Note [Local recursive groups]
  1023. ; let all_usg = spec_usg `combineUsage` body_usg -- Note [spec_usg includes rhs_usg]
  1024. bind' = Rec (concat (zipWith ruleInfoBinds rhs_infos specs))
  1025. ; return (all_usg { scu_calls = scu_calls all_usg `delVarEnvList` bndrs' },
  1026. Let bind' body') }
  1027. {-
  1028. Note [Local let bindings]
  1029. ~~~~~~~~~~~~~~~~~~~~~~~~~
  1030. It is not uncommon to find this
  1031. let $j = \x. <blah> in ...$j True...$j True...
  1032. Here $j is an arbitrary let-bound function, but it often comes up for
  1033. join points. We might like to specialise $j for its call patterns.
  1034. Notice the difference from a letrec, where we look for call patterns
  1035. in the *RHS* of the function. Here we look for call patterns in the
  1036. *body* of the let.
  1037. At one point I predicated this on the RHS mentioning the outer
  1038. recursive function, but that's not essential and might even be
  1039. harmful. I'm not sure.
  1040. -}
  1041. scApp :: ScEnv -> (InExpr, [InExpr]) -> UniqSM (ScUsage, CoreExpr)
  1042. scApp env (Var fn, args) -- Function is a variable
  1043. = ASSERT( not (null args) )
  1044. do { args_w_usgs <- mapM (scExpr env) args
  1045. ; let (arg_usgs, args') = unzip args_w_usgs
  1046. arg_usg = combineUsages arg_usgs
  1047. ; case scSubstId env fn of
  1048. fn'@(Lam {}) -> scExpr (zapScSubst env) (doBeta fn' args')
  1049. -- Do beta-reduction and try again
  1050. Var fn' -> return (arg_usg `combineUsage` mkVarUsage env fn' args',
  1051. mkApps (Var fn') args')
  1052. other_fn' -> return (arg_usg, mkApps other_fn' args') }
  1053. -- NB: doing this ignores any usage info from the substituted
  1054. -- function, but I don't think that matters. If it does
  1055. -- we can fix it.
  1056. where
  1057. doBeta :: OutExpr -> [OutExpr] -> OutExpr
  1058. -- ToDo: adjust for System IF
  1059. doBeta (Lam bndr body) (arg : args) = Let (NonRec bndr arg) (doBeta body args)
  1060. doBeta fn args = mkApps fn args
  1061. -- The function is almost always a variable, but not always.
  1062. -- In particular, if this pass follows float-in,
  1063. -- which it may, we can get
  1064. -- (let f = ...f... in f) arg1 arg2
  1065. scApp env (other_fn, args)
  1066. = do { (fn_usg, fn') <- scExpr env other_fn
  1067. ; (arg_usgs, args') <- mapAndUnzipM (scExpr env) args
  1068. ; return (combineUsages arg_usgs `combineUsage` fn_usg, mkApps fn' args') }
  1069. ----------------------
  1070. mkVarUsage :: ScEnv -> Id -> [CoreExpr] -> ScUsage
  1071. mkVarUsage env fn args
  1072. = case lookupHowBound env fn of
  1073. Just RecFun -> SCU { scu_calls = unitVarEnv fn [Call fn args (sc_vals env)]
  1074. , scu_occs = emptyVarEnv }
  1075. Just RecArg -> SCU { scu_calls = emptyVarEnv
  1076. , scu_occs = unitVarEnv fn arg_occ }
  1077. Nothing -> nullUsage
  1078. where
  1079. -- I rather think we could use UnkOcc all the time
  1080. arg_occ | null args = UnkOcc
  1081. | otherwise = evalScrutOcc
  1082. ----------------------
  1083. scTopBindEnv :: ScEnv -> CoreBind -> UniqSM (ScEnv, CoreBind)
  1084. scTopBindEnv env (Rec prs)
  1085. = do { let (rhs_env1,bndrs') = extendRecBndrs env bndrs
  1086. rhs_env2 = extendHowBound rhs_env1 bndrs RecFun
  1087. prs' = zip bndrs' rhss
  1088. ; return (rhs_env2, Rec prs') }
  1089. where
  1090. (bndrs,rhss) = unzip prs
  1091. scTopBindEnv env (NonRec bndr rhs)
  1092. = do { let (env1, bndr') = extendBndr env bndr
  1093. env2 = extendValEnv env1 bndr' (isValue (sc_vals env) rhs)
  1094. ; return (env2, NonRec bndr' rhs) }
  1095. ----------------------
  1096. scTopBind :: ScEnv -> ScUsage -> CoreBind -> UniqSM (ScUsage, CoreBind)
  1097. {-
  1098. scTopBind _ usage _
  1099. | pprTrace "scTopBind_usage" (ppr (scu_calls usage)) False
  1100. = error "false"
  1101. -}
  1102. scTopBind env body_usage (Rec prs)
  1103. | Just threshold <- sc_size env
  1104. , not force_spec
  1105. , not (all (couldBeSmallEnoughToInline (sc_dflags env) threshold) rhss)
  1106. -- No specialisation
  1107. = do { (rhs_usgs, rhss') <- mapAndUnzipM (scExpr env) rhss
  1108. ; return (body_usage `combineUsage` combineUsages rhs_usgs, Rec (bndrs `zip` rhss')) }
  1109. | otherwise -- Do specialisation
  1110. = do { rhs_infos <- mapM (scRecRhs env) prs
  1111. ; (spec_usage, specs) <- specRec TopLevel (scForce env force_spec)
  1112. body_usage rhs_infos
  1113. ; return (body_usage `combineUsage` spec_usage,
  1114. Rec (concat (zipWith ruleInfoBinds rhs_infos specs))) }
  1115. where
  1116. (bndrs,rhss) = unzip prs
  1117. force_spec = any (forceSpecBndr env) bndrs
  1118. -- Note [Forcing specialisation]
  1119. scTopBind env usage (NonRec bndr rhs) -- Oddly, we don't seem to specialise top-level non-rec functions
  1120. = do { (rhs_usg', rhs') <- scExpr env rhs
  1121. ; return (usage `combineUsage` rhs_usg', NonRec bndr rhs') }
  1122. ----------------------
  1123. scRecRhs :: ScEnv -> (OutId, InExpr) -> UniqSM RhsInfo
  1124. scRecRhs env (bndr,rhs)
  1125. = do { let (arg_bndrs,body) = collectBinders rhs
  1126. (body_env, arg_bndrs') = extendBndrsWith RecArg env arg_bndrs
  1127. ; (body_usg, body') <- scExpr body_env body
  1128. ; let (rhs_usg, arg_occs) = lookupOccs body_usg arg_bndrs'
  1129. ; return (RI { ri_rhs_usg = rhs_usg
  1130. , ri_fn = bndr, ri_new_rhs = mkLams arg_bndrs' body'
  1131. , ri_lam_bndrs = arg_bndrs, ri_lam_body = body
  1132. , ri_arg_occs = arg_occs }) }
  1133. -- The arg_occs says how the visible,
  1134. -- lambda-bound binders of the RHS are used
  1135. -- (including the TyVar binders)
  1136. -- Two pats are the same if they match both ways
  1137. ----------------------
  1138. ruleInfoBinds :: RhsInfo -> [OneSpec] -> [(Id,CoreExpr)]
  1139. ruleInfoBinds (RI { ri_fn = fn, ri_new_rhs = new_rhs }) specs
  1140. = [(id,rhs) | OS _ _ id rhs <- specs] ++
  1141. -- First the specialised bindings
  1142. [(fn `addIdSpecialisations` rules, new_rhs)]
  1143. -- And now the original binding
  1144. where
  1145. rules = [r | OS _ r _ _ <- specs]
  1146. {-
  1147. ************************************************************************
  1148. * *
  1149. The specialiser itself
  1150. * *
  1151. ************************************************************************
  1152. -}
  1153. data RhsInfo
  1154. = RI { ri_fn :: OutId -- The binder
  1155. , ri_new_rhs :: OutExpr -- The specialised RHS (in current envt)
  1156. , ri_rhs_usg :: ScUsage -- Usage info from specialising RHS
  1157. , ri_lam_bndrs :: [InVar] -- The *original* RHS (\xs.body)
  1158. , ri_lam_body :: InExpr -- Note [Specialise original body]
  1159. , ri_arg_occs :: [ArgOcc] -- Info on how the xs occur in body
  1160. }
  1161. data RuleInfo = SI [OneSpec] -- The specialisations we have generated
  1162. Int -- Length of specs; used for numbering them
  1163. (Maybe ScUsage) -- Just cs => we have not yet used calls in the
  1164. -- from calls in the *original* RHS as
  1165. -- seeds for new specialisations;
  1166. -- if you decide to do so, here is the
  1167. -- RHS usage (which has not yet been
  1168. -- unleashed)
  1169. -- Nothing => we have
  1170. -- See Note [Local recursive groups]
  1171. -- See Note [spec_usg includes rhs_usg]
  1172. -- One specialisation: Rule plus definition
  1173. data OneSpec = OS CallPat -- Call pattern that generated this specialisation
  1174. CoreRule -- Rule connecting original id with the specialisation
  1175. OutId OutExpr -- Spec id + its rhs
  1176. ----------------------
  1177. specNonRec :: ScEnv
  1178. -> ScUsage -- Body usage
  1179. -> RhsInfo -- Structure info usage info for un-specialised RHS
  1180. -> UniqSM (ScUsage, [OneSpec]) -- Usage from RHSs (specialised and not)
  1181. -- plus details of specialisations
  1182. specNonRec env body_usg rhs_info
  1183. = do { (spec_usg, SI specs _ _) <- specialise env (scu_calls body_usg)
  1184. rhs_info
  1185. (SI [] 0 (Just (ri_rhs_usg rhs_info)))
  1186. ; return (spec_usg, specs) }
  1187. ----------------------
  1188. specRec :: TopLevelFlag -> ScEnv
  1189. -> ScUsage -- Body usage
  1190. -> [RhsInfo] -- Structure info and usage info for un-specialised RHSs
  1191. -> UniqSM (ScUsage, [[OneSpec]]) -- Usage from all RHSs (specialised and not)
  1192. -- plus details of specialisations
  1193. specRec top_lvl env body_usg rhs_infos
  1194. = do { (spec_usg, spec_infos) <- go seed_calls nullUsage init_spec_infos
  1195. ; return (spec_usg, [ s | SI s _ _ <- spec_infos ]) }
  1196. where
  1197. (seed_calls, init_spec_infos) -- Note [Seeding top-level recursive groups]
  1198. | isTopLevel top_lvl
  1199. , any (isExportedId . ri_fn) rhs_infos -- Seed from body and RHSs
  1200. = (all_calls, [SI [] 0 Nothing | _ <- rhs_infos])
  1201. | otherwise -- Seed from body only
  1202. = (calls_in_body, [SI [] 0 (Just (ri_rhs_usg ri)) | ri <- rhs_infos])
  1203. calls_in_body = scu_calls body_usg
  1204. calls_in_rhss = foldr (combineCalls . scu_calls . ri_rhs_usg) emptyVarEnv rhs_infos
  1205. all_calls = calls_in_rhss `combineCalls` calls_in_body
  1206. -- Loop, specialising, until you get no new specialisations
  1207. go seed_calls usg_so_far spec_infos
  1208. | isEmptyVarEnv seed_calls
  1209. = return (usg_so_far, spec_infos)
  1210. | otherwise
  1211. = do { specs_w_usg <- zipWithM (specialise env seed_calls) rhs_infos spec_infos
  1212. ; let (extra_usg_s, new_spec_infos) = unzip specs_w_usg
  1213. extra_usg = combineUsages extra_usg_s
  1214. all_usg = usg_so_far `combineUsage` extra_usg
  1215. ; go (scu_calls extra_usg) all_usg new_spec_infos }
  1216. ----------------------
  1217. specialise
  1218. :: ScEnv
  1219. -> CallEnv -- Info on newly-discovered calls to this function
  1220. -> RhsInfo
  1221. -> RuleInfo -- Original RHS plus patterns dealt with
  1222. -> UniqSM (ScUsage, RuleInfo) -- New specialised versions and their usage
  1223. -- See Note [spec_usg includes rhs_usg]
  1224. -- Note: this only generates *specialised* bindings
  1225. -- The original binding is added by ruleInfoBinds
  1226. --
  1227. -- Note: the rhs here is the optimised version of the original rhs
  1228. -- So when we make a specialised copy of the RHS, we're starting
  1229. -- from an RHS whose nested functions have been optimised already.
  1230. specialise env bind_calls (RI { ri_fn = fn, ri_lam_bndrs = arg_bndrs
  1231. , ri_lam_body = body, ri_arg_occs = arg_occs })
  1232. spec_info@(SI specs spec_count mb_unspec)
  1233. | isBottomingId fn -- Note [Do not specialise diverging functions]
  1234. -- and do not generate specialisation seeds from its RHS
  1235. = return (nullUsage, spec_info)
  1236. | isNeverActive (idInlineActivation fn) -- See Note [Transfer activation]
  1237. || null arg_bndrs -- Only specialise functions
  1238. = case mb_unspec of -- Behave as if there was a single, boring call
  1239. Just rhs_usg -> return (rhs_usg, SI specs spec_count Nothing)
  1240. -- See Note [spec_usg includes rhs_usg]
  1241. Nothing -> return (nullUsage, spec_info)
  1242. | Just all_calls <- lookupVarEnv bind_calls fn
  1243. = -- pprTrace "specialise entry {" (ppr fn <+> ppr (length all_calls)) $
  1244. do { (boring_call, pats) <- callsToPats env specs arg_occs all_calls
  1245. -- Bale out if too many specialisations
  1246. ; let n_pats = length pats
  1247. spec_count' = n_pats + spec_count
  1248. ; case sc_count env of
  1249. Just max | not (sc_force env) && spec_count' > max
  1250. -> if (debugIsOn || opt_PprStyle_Debug) -- Suppress this scary message for
  1251. then pprTrace "SpecConstr" msg $ -- ordinary users! Trac #5125
  1252. return (nullUsage, spec_info)
  1253. else return (nullUsage, spec_info)
  1254. where
  1255. msg = vcat [ sep [ text "Function" <+> quotes (ppr fn)
  1256. , nest 2 (text "has" <+>
  1257. speakNOf spec_count' (text "call pattern") <> comma <+>
  1258. text "but the limit is" <+> int max) ]
  1259. , text "Use -fspec-constr-count=n to set the bound"
  1260. , extra ]
  1261. extra | not opt_PprStyle_Debug = text "Use -dppr-debug to see specialisations"
  1262. | otherwise = text "Specialisations:" <+> ppr (pats ++ [p | OS p _ _ _ <- specs])
  1263. _normal_case -> do {
  1264. -- ; if (not (null pats) || isJust mb_unspec) then
  1265. -- pprTrace "specialise" (vcat [ ppr fn <+> text "with" <+> int (length pats) <+> text "good patterns"
  1266. -- , text "mb_unspec" <+> ppr (isJust mb_unspec)
  1267. -- , text "arg_occs" <+> ppr arg_occs
  1268. -- , text "good pats" <+> ppr pats]) $
  1269. -- return ()
  1270. -- else return ()
  1271. ; let spec_env = decreaseSpecCount env n_pats
  1272. ; (spec_usgs, new_specs) <- mapAndUnzipM (spec_one spec_env fn arg_bndrs body)
  1273. (pats `zip` [spec_count..])
  1274. -- See Note [Specialise original body]
  1275. ; let spec_usg = combineUsages spec_usgs
  1276. -- If there were any boring calls among the seeds (= all_calls), then those
  1277. -- calls will call the un-specialised function. So we should use the seeds
  1278. -- from the _unspecialised_ function's RHS, which are in mb_unspec, by returning
  1279. -- then in new_usg.
  1280. (new_usg, mb_unspec')
  1281. = case mb_unspec of
  1282. Just rhs_usg | boring_call -> (spec_usg `combineUsage` rhs_usg, Nothing)
  1283. _ -> (spec_usg, mb_unspec)
  1284. -- ; pprTrace "specialise return }" (ppr fn
  1285. -- <+> ppr (scu_calls new_usg))
  1286. ; return (new_usg, SI (new_specs ++ specs) spec_count' mb_unspec') } }
  1287. | otherwise -- No new seeds, so return nullUsage
  1288. = return (nullUsage, spec_info)
  1289. ---------------------
  1290. spec_one :: ScEnv
  1291. -> OutId -- Function
  1292. -> [InVar] -- Lambda-binders of RHS; should match patterns
  1293. -> InExpr -- Body of the original function
  1294. -> (CallPat, Int)
  1295. -> UniqSM (ScUsage, OneSpec) -- Rule and binding
  1296. -- spec_one creates a specialised copy of the function, together
  1297. -- with a rule for using it. I'm very proud of how short this
  1298. -- function is, considering what it does :-).
  1299. {-
  1300. Example
  1301. In-scope: a, x::a
  1302. f = /\b \y::[(a,b)] -> ....f (b,c) ((:) (a,(b,c)) (x,v) (h w))...
  1303. [c::*, v::(b,c) are presumably bound by the (...) part]
  1304. ==>
  1305. f_spec = /\ b c \ v::(b,c) hw::[(a,(b,c))] ->
  1306. (...entire body of f...) [b -> (b,c),
  1307. y -> ((:) (a,(b,c)) (x,v) hw)]
  1308. RULE: forall b::* c::*, -- Note, *not* forall a, x
  1309. v::(b,c),
  1310. hw::[(a,(b,c))] .
  1311. f (b,c) ((:) (a,(b,c)) (x,v) hw) = f_spec b c v hw
  1312. -}
  1313. spec_one env fn arg_bndrs body (call_pat@(qvars, pats), rule_number)
  1314. = do { spec_uniq <- getUniqueM
  1315. ; let spec_env = extendScSubstList (extendScInScope env qvars)
  1316. (arg_bndrs `zip` pats)
  1317. fn_name = idName fn
  1318. fn_loc = nameSrcSpan fn_name
  1319. fn_occ = nameOccName fn_name
  1320. spec_occ = mkSpecOcc fn_occ
  1321. -- We use fn_occ rather than fn in the rule_name string
  1322. -- as we don't want the uniq to end up in the rule, and
  1323. -- hence in the ABI, as that can cause spurious ABI
  1324. -- changes (#4012).
  1325. rule_name = mkFastString ("SC:" ++ occNameString fn_occ ++ show rule_number)
  1326. spec_name = mkInternalName spec_uniq spec_occ fn_loc
  1327. -- ; pprTrace "{spec_one" (ppr (sc_count env) <+> ppr fn <+> ppr pats <+> text "-->" <+> ppr spec_name) $
  1328. -- return ()
  1329. -- Specialise the body
  1330. ; (spec_usg, spec_body) <- scExpr spec_env body
  1331. -- ; pprTrace "done spec_one}" (ppr fn) $
  1332. -- return ()
  1333. -- And build the results
  1334. ; let spec_id = mkLocalIdOrCoVar spec_name (mkLamTypes spec_lam_args body_ty)
  1335. -- See Note [Transfer strictness]
  1336. `setIdStrictness` spec_str
  1337. `setIdArity` count isId spec_lam_args
  1338. spec_str = calcSpecStrictness fn spec_lam_args pats
  1339. -- Conditionally use result of new worker-wrapper transform
  1340. (spec_lam_args, spec_call_args) = mkWorkerArgs (sc_dflags env) qvars body_ty
  1341. -- Usual w/w hack to avoid generating
  1342. -- a spec_rhs of unlifted type and no args
  1343. spec_lam_args_str = handOutStrictnessInformation (fst (splitStrictSig spec_str)) spec_lam_args
  1344. -- Annotate the variables with the strictness information from
  1345. -- the function (see Note [Strictness information in worker binders])
  1346. spec_rhs = mkLams spec_lam_args_str spec_body
  1347. body_ty = exprType spec_body
  1348. rule_rhs = mkVarApps (Var spec_id) spec_call_args
  1349. inline_act = idInlineActivation fn
  1350. this_mod = sc_module spec_env
  1351. rule = mkRule this_mod True {- Auto -} True {- Local -}
  1352. rule_name inline_act fn_name qvars pats rule_rhs
  1353. -- See Note [Transfer activation]
  1354. ; return (spec_usg, OS call_pat rule spec_id spec_rhs) }
  1355. -- See Note [Strictness information in worker binders]
  1356. handOutStrictnessInformation :: [Demand] -> [Var] -> [Var]
  1357. handOutStrictnessInformation = go
  1358. where
  1359. go _ [] = []
  1360. go [] vs = vs
  1361. go (d:dmds) (v:vs) | isId v = setIdDemandInfo v d : go dmds vs
  1362. go dmds (v:vs) = v : go dmds vs
  1363. calcSpecStrictness :: Id -- The original function
  1364. -> [Var] -> [CoreExpr] -- Call pattern
  1365. -> StrictSig -- Strictness of specialised thing
  1366. -- See Note [Transfer strictness]
  1367. calcSpecStrictness fn qvars pats
  1368. = mkClosedStrictSig spec_dmds topRes
  1369. where
  1370. spec_dmds = [ lookupVarEnv dmd_env qv `orElse` topDmd | qv <- qvars, isId qv ]
  1371. StrictSig (DmdType _ dmds _) = idStrictness fn
  1372. dmd_env = go emptyVarEnv dmds pats
  1373. go :: DmdEnv -> [Demand] -> [CoreExpr] -> DmdEnv
  1374. go env ds (Type {} : pats) = go env ds pats
  1375. go env ds (Coercion {} : pats) = go env ds pats
  1376. go env (d:ds) (pat : pats) = go (go_one env d pat) ds pats
  1377. go env _ _ = env
  1378. go_one :: DmdEnv -> Demand -> CoreExpr -> DmdEnv
  1379. go_one env d (Var v) = extendVarEnv_C bothDmd env v d
  1380. go_one env d e
  1381. | Just ds <- splitProdDmd_maybe d -- NB: d does not have to be strict
  1382. , (Var _, args) <- collectArgs e = go env ds args
  1383. go_one env _ _ = env
  1384. {-
  1385. Note [spec_usg includes rhs_usg]
  1386. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1387. In calls to 'specialise', the returned ScUsage must include the rhs_usg in
  1388. the passed-in RuleInfo, unless there are no calls at all to the function.
  1389. The caller can, indeed must, assume this. He should not combine in rhs_usg
  1390. himself, or he'll get rhs_usg twice -- and that can lead to an exponential
  1391. blowup of duplicates in the CallEnv. This is what gave rise to the massive
  1392. performace loss in Trac #8852.
  1393. Note [Specialise original body]
  1394. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1395. The RhsInfo for a binding keeps the *original* body of the binding. We
  1396. must specialise that, *not* the result of applying specExpr to the RHS
  1397. (which is also kept in RhsInfo). Otherwise we end up specialising a
  1398. specialised RHS, and that can lead directly to exponential behaviour.
  1399. Note [Transfer activation]
  1400. ~~~~~~~~~~~~~~~~~~~~~~~~~~
  1401. This note is for SpecConstr, but exactly the same thing
  1402. happens in the overloading specialiser; see
  1403. Note [Auto-specialisation and RULES] in Specialise.
  1404. In which phase should the specialise-constructor rules be active?
  1405. Originally I made them always-active, but Manuel found that this
  1406. defeated some clever user-written rules. Then I made them active only
  1407. in Phase 0; after all, currently, the specConstr transformation is
  1408. only run after the simplifier has reached Phase 0, but that meant
  1409. that specialisations didn't fire inside wrappers; see test
  1410. simplCore/should_compile/spec-inline.
  1411. So now I just use the inline-activation of the parent Id, as the
  1412. activation for the specialiation RULE, just like the main specialiser;
  1413. This in turn means there is no point in specialising NOINLINE things,
  1414. so we test for that.
  1415. Note [Transfer strictness]
  1416. ~~~~~~~~~~~~~~~~~~~~~~~~~~
  1417. We must transfer strictness information from the original function to
  1418. the specialised one. Suppose, for example
  1419. f has strictness SS
  1420. and a RULE f (a:as) b = f_spec a as b
  1421. Now we want f_spec to have strictness LLS, otherwise we'll use call-by-need
  1422. when calling f_spec instead of call-by-value. And that can result in
  1423. unbounded worsening in space (cf the classic foldl vs foldl')
  1424. See Trac #3437 for a good example.
  1425. The function calcSpecStrictness performs the calculation.
  1426. Note [Strictness information in worker binders]
  1427. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1428. After having calculated the strictness annotation for the worker (see Note
  1429. [Transfer strictness] above), we also want to have this information attached to
  1430. the worker’s arguments, for the benefit of later passes. The function
  1431. handOutStrictnessInformation decomposes the strictness annotation calculated by
  1432. calcSpecStrictness and attaches them to the variables.
  1433. ************************************************************************
  1434. * *
  1435. \subsection{Argument analysis}
  1436. * *
  1437. ************************************************************************
  1438. This code deals with analysing call-site arguments to see whether
  1439. they are constructor applications.
  1440. Note [Free type variables of the qvar types]
  1441. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1442. In a call (f @a x True), that we want to specialise, what variables should
  1443. we quantify over. Clearly over 'a' and 'x', but what about any type variables
  1444. free in x's type? In fact we don't need to worry about them because (f @a)
  1445. can only be a well-typed application if its type is compatible with x, so any
  1446. variables free in x's type must be free in (f @a), and hence either be gathered
  1447. via 'a' itself, or be in scope at f's defn. Hence we just take
  1448. (exprsFreeVars pats).
  1449. BUT phantom type synonyms can mess this reasoning up,
  1450. eg x::T b with type T b = Int
  1451. So we apply expandTypeSynonyms to the bound Ids.
  1452. See Trac # 5458. Yuk.
  1453. Note [SpecConstr call patterns]
  1454. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1455. A "call patterns" that we collect is going to become the LHS of a RULE.
  1456. It's important that it doesn't have
  1457. e |> Refl
  1458. or
  1459. e |> g1 |> g2
  1460. because both of these will be optimised by Simplify.simplRule. In the
  1461. former case such optimisation benign, because the rule will match more
  1462. terms; but in the latter we may lose a binding of 'g1' or 'g2', and
  1463. end up with a rule LHS that doesn't bind the template variables
  1464. (Trac #10602).
  1465. The simplifier eliminates such things, but SpecConstr itself constructs
  1466. new terms by substituting. So the 'mkCast' in the Cast case of scExpr
  1467. is very important!
  1468. -}
  1469. type CallPat = ([Var], [CoreExpr]) -- Quantified variables and arguments
  1470. -- See Note [SpecConstr call patterns]
  1471. callsToPats :: ScEnv -> [OneSpec] -> [ArgOcc] -> [Call] -> UniqSM (Bool, [CallPat])
  1472. -- Result has no duplicate patterns,
  1473. -- nor ones mentioned in done_pats
  1474. -- Bool indicates that there was at least one boring pattern
  1475. callsToPats env done_specs bndr_occs calls
  1476. = do { mb_pats <- mapM (callToPats env bndr_occs) calls
  1477. ; let good_pats :: [(CallPat, ValueEnv)]
  1478. good_pats = catMaybes mb_pats
  1479. done_pats = [p | OS p _ _ _ <- done_specs]
  1480. is_done p = any (samePat p) done_pats
  1481. no_recursive = map fst (filterOut (is_too_recursive env) good_pats)
  1482. ; return (any isNothing mb_pats,
  1483. filterOut is_done (nubBy samePat no_recursive)) }
  1484. is_too_recursive :: ScEnv -> (CallPat, ValueEnv) -> Bool
  1485. -- Count the number of recursive constructors in a call pattern,
  1486. -- filter out if there are more than the maximum.
  1487. -- This is only necessary if ForceSpecConstr is in effect:
  1488. -- otherwise specConstrCount will cause specialisation to terminate.
  1489. -- See Note [Limit recursive specialisation]
  1490. -- TODO: make me more accurate
  1491. is_too_recursive env ((_,exprs), val_env)
  1492. = sc_force env && maximum (map go exprs) > sc_recursive env
  1493. where
  1494. go e
  1495. | Just (ConVal (DataAlt _) args) <- isValue val_env e
  1496. = 1 + sum (map go args)
  1497. | App f a <- e
  1498. = go f + go a
  1499. | otherwise
  1500. = 0
  1501. callToPats :: ScEnv -> [ArgOcc] -> Call -> UniqSM (Maybe (CallPat, ValueEnv))
  1502. -- The [Var] is the variables to quantify over in the rule
  1503. -- Type variables come first, since they may scope
  1504. -- over the following term variables
  1505. -- The [CoreExpr] are the argument patterns for the rule
  1506. callToPats env bndr_occs (Call _ args con_env)
  1507. | length args < length bndr_occs -- Check saturated
  1508. = return Nothing
  1509. | otherwise
  1510. = do { let in_scope = substInScope (sc_subst env)
  1511. ; (interesting, pats) <- argsToPats env in_scope con_env args bndr_occs
  1512. ; let pat_fvs = exprsFreeVarsList pats
  1513. -- To get determinism we need the list of free variables in
  1514. -- deterministic order. Otherwise we end up creating
  1515. -- lambdas with different argument orders. See
  1516. -- determinism/simplCore/should_compile/spec-inline-determ.hs
  1517. -- for an example. For explanation of determinism
  1518. -- considerations See Note [Unique Determinism] in Unique.
  1519. in_scope_vars = getInScopeVars in_scope
  1520. qvars = filterOut (`elemVarSet` in_scope_vars) pat_fvs
  1521. -- Quantify over variables that are not in scope
  1522. -- at the call site
  1523. -- See Note [Free type variables of the qvar types]
  1524. -- See Note [Shadowing] at the top
  1525. (ktvs, ids) = partition isTyVar qvars
  1526. qvars' = toposortTyVars ktvs ++ map sanitise ids
  1527. -- Order into kind variables, type variables, term variables
  1528. -- The kind of a type variable may mention a kind variable
  1529. -- and the type of a term variable may mention a type variable
  1530. sanitise id = id `setIdType` expandTypeSynonyms (idType id)
  1531. -- See Note [Free type variables of the qvar types]
  1532. ; -- pprTrace "callToPats" (ppr args $$ ppr bndr_occs) $
  1533. if interesting
  1534. then return (Just ((qvars', pats), con_env))
  1535. else return Nothing }
  1536. -- argToPat takes an actual argument, and returns an abstracted
  1537. -- version, consisting of just the "constructor skeleton" of the
  1538. -- argument, with non-constructor sub-expression replaced by new
  1539. -- placeholder variables. For example:
  1540. -- C a (D (f x) (g y)) ==> C p1 (D p2 p3)
  1541. argToPat :: ScEnv
  1542. -> InScopeSet -- What's in scope at the fn defn site
  1543. -> ValueEnv -- ValueEnv at the call site
  1544. -> CoreArg -- A call arg (or component thereof)
  1545. -> ArgOcc
  1546. -> UniqSM (Bool, CoreArg)
  1547. -- Returns (interesting, pat),
  1548. -- where pat is the pattern derived from the argument
  1549. -- interesting=True if the pattern is non-trivial (not a variable or type)
  1550. -- E.g. x:xs --> (True, x:xs)
  1551. -- f xs --> (False, w) where w is a fresh wildcard
  1552. -- (f xs, 'c') --> (True, (w, 'c')) where w is a fresh wildcard
  1553. -- \x. x+y --> (True, \x. x+y)
  1554. -- lvl7 --> (True, lvl7) if lvl7 is bound
  1555. -- somewhere further out
  1556. argToPat _env _in_scope _val_env arg@(Type {}) _arg_occ
  1557. = return (False, arg)
  1558. argToPat env in_scope val_env (Tick _ arg) arg_occ
  1559. = argToPat env in_scope val_env arg arg_occ
  1560. -- Note [Notes in call patterns]
  1561. -- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1562. -- Ignore Notes. In particular, we want to ignore any InlineMe notes
  1563. -- Perhaps we should not ignore profiling notes, but I'm going to
  1564. -- ride roughshod over them all for now.
  1565. --- See Note [Notes in RULE matching] in Rules
  1566. argToPat env in_scope val_env (Let _ arg) arg_occ
  1567. = argToPat env in_scope val_env arg arg_occ
  1568. -- See Note [Matching lets] in Rule.hs
  1569. -- Look through let expressions
  1570. -- e.g. f (let v = rhs in (v,w))
  1571. -- Here we can specialise for f (v,w)
  1572. -- because the rule-matcher will look through the let.
  1573. {- Disabled; see Note [Matching cases] in Rule.hs
  1574. argToPat env in_scope val_env (Case scrut _ _ [(_, _, rhs)]) arg_occ
  1575. | exprOkForSpeculation scrut -- See Note [Matching cases] in Rule.hhs
  1576. = argToPat env in_scope val_env rhs arg_occ
  1577. -}
  1578. argToPat env in_scope val_env (Cast arg co) arg_occ
  1579. | not (ignoreType env ty2)
  1580. = do { (interesting, arg') <- argToPat env in_scope val_env arg arg_occ
  1581. ; if not interesting then
  1582. wildCardPat ty2
  1583. else do
  1584. { -- Make a wild-card pattern for the coercion
  1585. uniq <- getUniqueM
  1586. ; let co_name = mkSysTvName uniq (fsLit "sg")
  1587. co_var = mkCoVar co_name (mkCoercionType Representational ty1 ty2)
  1588. ; return (interesting, Cast arg' (mkCoVarCo co_var)) } }
  1589. where
  1590. Pair ty1 ty2 = coercionKind co
  1591. {- Disabling lambda specialisation for now
  1592. It's fragile, and the spec_loop can be infinite
  1593. argToPat in_scope val_env arg arg_occ
  1594. | is_value_lam arg
  1595. = return (True, arg)
  1596. where
  1597. is_value_lam (Lam v e) -- Spot a value lambda, even if
  1598. | isId v = True -- it is inside a type lambda
  1599. | otherwise = is_value_lam e
  1600. is_value_lam other = False
  1601. -}
  1602. -- Check for a constructor application
  1603. -- NB: this *precedes* the Var case, so that we catch nullary constrs
  1604. argToPat env in_scope val_env arg arg_occ
  1605. | Just (ConVal (DataAlt dc) args) <- isValue val_env arg
  1606. , not (ignoreDataCon env dc) -- See Note [NoSpecConstr]
  1607. , Just arg_occs <- mb_scrut dc
  1608. = do { let (ty_args, rest_args) = splitAtList (dataConUnivTyVars dc) args
  1609. ; (_, args') <- argsToPats env in_scope val_env rest_args arg_occs
  1610. ; return (True,
  1611. mkConApp dc (ty_args ++ args')) }
  1612. where
  1613. mb_scrut dc = case arg_occ of
  1614. ScrutOcc bs
  1615. | Just occs <- lookupUFM bs dc
  1616. -> Just (occs) -- See Note [Reboxing]
  1617. _other | sc_force env -> Just (repeat UnkOcc)
  1618. | otherwise -> Nothing
  1619. -- Check if the argument is a variable that
  1620. -- (a) is used in an interesting way in the function body
  1621. -- (b) we know what its value is
  1622. -- In that case it counts as "interesting"
  1623. argToPat env in_scope val_env (Var v) arg_occ
  1624. | sc_force env || case arg_occ of { UnkOcc -> False; _other -> True }, -- (a)
  1625. is_value, -- (b)
  1626. not (ignoreType env (varType v))
  1627. = return (True, Var v)
  1628. where
  1629. is_value
  1630. | isLocalId v = v `elemInScopeSet` in_scope
  1631. && isJust (lookupVarEnv val_env v)
  1632. -- Local variables have values in val_env
  1633. | otherwise = isValueUnfolding (idUnfolding v)
  1634. -- Imports have unfoldings
  1635. -- I'm really not sure what this comment means
  1636. -- And by not wild-carding we tend to get forall'd
  1637. -- variables that are in scope, which in turn can
  1638. -- expose the weakness in let-matching
  1639. -- See Note [Matching lets] in Rules
  1640. -- Check for a variable bound inside the function.
  1641. -- Don't make a wild-card, because we may usefully share
  1642. -- e.g. f a = let x = ... in f (x,x)
  1643. -- NB: this case follows the lambda and con-app cases!!
  1644. -- argToPat _in_scope _val_env (Var v) _arg_occ
  1645. -- = return (False, Var v)
  1646. -- SLPJ : disabling this to avoid proliferation of versions
  1647. -- also works badly when thinking about seeding the loop
  1648. -- from the body of the let
  1649. -- f x y = letrec g z = ... in g (x,y)
  1650. -- We don't want to specialise for that *particular* x,y
  1651. -- The default case: make a wild-card
  1652. -- We use this for coercions too
  1653. argToPat _env _in_scope _val_env arg _arg_occ
  1654. = wildCardPat (exprType arg)
  1655. wildCardPat :: Type -> UniqSM (Bool, CoreArg)
  1656. wildCardPat ty
  1657. = do { uniq <- getUniqueM
  1658. ; let id = mkSysLocalOrCoVar (fsLit "sc") uniq ty
  1659. ; return (False, varToCoreExpr id) }
  1660. argsToPats :: ScEnv -> InScopeSet -> ValueEnv
  1661. -> [CoreArg] -> [ArgOcc] -- Should be same length
  1662. -> UniqSM (Bool, [CoreArg])
  1663. argsToPats env in_scope val_env args occs
  1664. = do { stuff <- zipWithM (argToPat env in_scope val_env) args occs
  1665. ; let (interesting_s, args') = unzip stuff
  1666. ; return (or interesting_s, args') }
  1667. isValue :: ValueEnv -> CoreExpr -> Maybe Value
  1668. isValue _env (Lit lit)
  1669. | litIsLifted lit = Nothing
  1670. | otherwise = Just (ConVal (LitAlt lit) [])
  1671. isValue env (Var v)
  1672. | Just cval <- lookupVarEnv env v
  1673. = Just cval -- You might think we could look in the idUnfolding here
  1674. -- but that doesn't take account of which branch of a
  1675. -- case we are in, which is the whole point
  1676. | not (isLocalId v) && isCheapUnfolding unf
  1677. = isValue env (unfoldingTemplate unf)
  1678. where
  1679. unf = idUnfolding v
  1680. -- However we do want to consult the unfolding
  1681. -- as well, for let-bound constructors!
  1682. isValue env (Lam b e)
  1683. | isTyVar b = case isValue env e of
  1684. Just _ -> Just LambdaVal
  1685. Nothing -> Nothing
  1686. | otherwise = Just LambdaVal
  1687. isValue env (Tick t e)
  1688. | not (tickishIsCode t)
  1689. = isValue env e
  1690. isValue _env expr -- Maybe it's a constructor application
  1691. | (Var fun, args, _) <- collectArgsTicks (not . tickishIsCode) expr
  1692. = case isDataConWorkId_maybe fun of
  1693. Just con | args `lengthAtLeast` dataConRepArity con
  1694. -- Check saturated; might be > because the
  1695. -- arity excludes type args
  1696. -> Just (ConVal (DataAlt con) args)
  1697. _other | valArgCount args < idArity fun
  1698. -- Under-applied function
  1699. -> Just LambdaVal -- Partial application
  1700. _other -> Nothing
  1701. isValue _env _expr = Nothing
  1702. valueIsWorkFree :: Value -> Bool
  1703. valueIsWorkFree LambdaVal = True
  1704. valueIsWorkFree (ConVal _ args) = all exprIsWorkFree args
  1705. samePat :: CallPat -> CallPat -> Bool
  1706. samePat (vs1, as1) (vs2, as2)
  1707. = all2 same as1 as2
  1708. where
  1709. same (Var v1) (Var v2)
  1710. | v1 `elem` vs1 = v2 `elem` vs2
  1711. | v2 `elem` vs2 = False
  1712. | otherwise = v1 == v2
  1713. same (Lit l1) (Lit l2) = l1==l2
  1714. same (App f1 a1) (App f2 a2) = same f1 f2 && same a1 a2
  1715. same (Type {}) (Type {}) = True -- Note [Ignore type differences]
  1716. same (Coercion {}) (Coercion {}) = True
  1717. same (Tick _ e1) e2 = same e1 e2 -- Ignore casts and notes
  1718. same (Cast e1 _) e2 = same e1 e2
  1719. same e1 (Tick _ e2) = same e1 e2
  1720. same e1 (Cast e2 _) = same e1 e2
  1721. same e1 e2 = WARN( bad e1 || bad e2, ppr e1 $$ ppr e2)
  1722. False -- Let, lambda, case should not occur
  1723. bad (Case {}) = True
  1724. bad (Let {}) = True
  1725. bad (Lam {}) = True
  1726. bad _other = False
  1727. {-
  1728. Note [Ignore type differences]
  1729. ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  1730. We do not want to generate specialisations where the call patterns
  1731. differ only in their type arguments! Not only is it utterly useless,
  1732. but it also means that (with polymorphic recursion) we can generate
  1733. an infinite number of specialisations. Example is Data.Sequence.adjustTree,
  1734. I think.
  1735. -}