- -- process each binding
- (new_binds, g, srt, new_off) = doBinds binds emptyUniqSet [] off []
-
- doBinds [] g srt off new_binds = (reverse new_binds, g, srt, off)
- doBinds ((binder,rhs):binds) g srt off new_binds =
- srtRhs rho cont_refs off rhs =: \(rhs, g', srt', off) ->
- doBinds binds (unionUniqSets g g') (srt'++srt) off
- ((binder,rhs):new_binds)
-\end{code}
-
------------------------------------------------------------------------------
-Right Hand Sides
-
-\begin{code}
-srtRhs :: UniqFM CafInfo -> (UniqSet Id, UniqFM (UniqSet Id))
- -> Int -> StgRhs -> (StgRhs, UniqSet Id, [Id], Int)
-
-srtRhs rho cont off (StgRhsClosure cc bi old_srt free_vars u args body) =
- srtExpr rho cont off body =: \(body, g, srt, off) ->
- (StgRhsClosure cc bi old_srt free_vars u args body, g, srt, off)
-
-srtRhs rho cont off e@(StgRhsCon cc con args) =
- (e, getGlobalRefs rho args, [], off)
-\end{code}
-
------------------------------------------------------------------------------
-Expressions
-
-\begin{code}
-srtExpr :: UniqFM CafInfo -> (UniqSet Id, UniqFM (UniqSet Id))
- -> Int -> StgExpr -> (StgExpr, UniqSet Id, [Id], Int)
-
-srtExpr rho (cont,lne) off e@(StgApp f args) = (e, global_refs, [], off)
- where global_refs =
- cont `unionUniqSets`
- getGlobalRefs rho (StgVarArg f:args) `unionUniqSets`
- lookupPossibleLNE lne f
-
-srtExpr rho (cont,lne) off e@(StgLit l) = (e, cont, [], off)
-
-srtExpr rho (cont,lne) off e@(StgConApp con args) =
- (e, cont `unionUniqSets` getGlobalRefs rho args, [], off)
-
-srtExpr rho (cont,lne) off e@(StgPrimApp op args ty) =
- (e, cont `unionUniqSets` getGlobalRefs rho args, [], off)
-
-srtExpr rho c@(cont,lne) off (StgCase scrut live1 live2 uniq _{-srt-} alts) =
- srtCaseAlts rho c off alts =: \(alts, alts_g, alts_srt, alts_off) ->
-
- -- construct the SRT for this case
- let (this_srt, scrut_off) = construct_srt rho alts_g alts_srt alts_off in
-
- -- global refs in the continuation is alts_g.
- srtExpr rho (alts_g,lne) scrut_off scrut
- =: \(scrut, scrut_g, scrut_srt, case_off) ->
- let
- g = unionUniqSets alts_g scrut_g
- srt = scrut_srt ++ this_srt
- srt_info = case length this_srt of
- 0 -> NoSRT
- len -> SRT off len
- in
- (StgCase scrut live1 live2 uniq srt_info alts, g, srt, case_off)
-
-srtExpr rho cont off (StgLet bind body) =
- srtLet rho cont off bind body StgLet (\_ cont -> cont)
-
-srtExpr rho cont off (StgLetNoEscape live1 live2 b@(StgNonRec bndr rhs) body)
- = srtLet rho cont off b body (StgLetNoEscape live1 live2) calc_cont
- where calc_cont g (cont,lne) = (cont,addToUFM lne bndr g)
-
--- for recursive let-no-escapes, we do *two* passes, the first time
--- just to extract the list of global refs, and the second time we actually
--- construct the SRT now that we know what global refs should be in
--- the various let-no-escape continuations.
-srtExpr rho conts@(cont,lne) off
- (StgLetNoEscape live1 live2 bind@(StgRec pairs) body)
- = srtBind rho conts off bind =: \(_, g, _, _) ->
- let
- lne' = addListToUFM lne [ (bndr,g) | (bndr,_) <- pairs ]
- calc_cont _ conts = conts
- in
- srtLet rho (cont,lne') off bind body (StgLetNoEscape live1 live2) calc_cont
-
-
-srtExpr rho cont off (StgSCC cc expr) =
- srtExpr rho cont off expr =: \(expr, g, srt, off) ->
- (StgSCC cc expr, g, srt, off)
-
-#ifdef DEBUG
-srtExpr rho cont off expr = pprPanic "srtExpr" (ppr expr)
-#else
-srtExpr rho cont off expr = panic "srtExpr"
-#endif
-\end{code}
-
------------------------------------------------------------------------------
-Let-expressions
-
-This is quite complicated stuff...
-
-\begin{code}
-srtLet rho cont off bind body let_constr calc_cont
-
- -- If the bindings are all constructors, then we don't need to
- -- buid an SRT at all...
- | all_con_binds bind =
- srtBind rho cont off bind =: \(bind, bind_g, bind_srt, off) ->
- srtExpr rho cont off body =: \(body, body_g, body_srt, off) ->
- let
- g = unionUniqSets bind_g body_g
- srt = body_srt ++ bind_srt
- in
- (let_constr bind body, g, srt, off)
-
- -- we have some closure bindings...
- | otherwise =
-
- -- first, find the sub-SRTs in the binding
- srtBind rho cont off bind =: \(bind, bind_g, bind_srt, bind_off) ->
-
- -- construct the SRT for this binding
- let (this_srt, body_off) = construct_srt rho bind_g bind_srt bind_off in
-
- -- get the new continuation information (if a let-no-escape)
- let new_cont = calc_cont bind_g cont in
-
- -- now find the SRTs in the body
- srtExpr rho new_cont body_off body =: \(body, body_g, body_srt, let_off) ->
-
- let
- -- union all the global references together
- let_g = unionUniqSets bind_g body_g
-
- -- concatenate the sub-SRTs
- let_srt = body_srt ++ this_srt
-
- -- attach the SRT info to the binding
- bind' = attach_srt_bind bind off (length this_srt)
- in
- (let_constr bind' body, let_g, let_srt, let_off)
-\end{code}
-
------------------------------------------------------------------------------
-Construct an SRT.
-
-Construct the SRT at this point from its sub-SRTs and any new global
-references which aren't already contained in one of the sub-SRTs (and
-which are "live").
-
-\begin{code}
-construct_srt rho global_refs sub_srt current_offset
- = let
- extra_refs = filter (`notElem` sub_srt)
- (filter (mayHaveCafRefs rho) (uniqSetToList global_refs))
- this_srt = extra_refs ++ sub_srt
-
- -- Add the length of the new entries to the
- -- current offset to get the next free offset in the global SRT.
- new_offset = current_offset + length extra_refs
- in (this_srt, new_offset)
-\end{code}
-
------------------------------------------------------------------------------
-Case Alternatives
+ (rhs', srt) = srtTopRhs b rhs
+ env' = maybeExtendEnv env b rhs
+ srt' = applyEnvList env srt
+srtTopBinds env (StgRec bs : binds) =
+ (StgRec (zip bndrs rhss), zip bndrs srts') : srtTopBinds env binds
+ where
+ (rhss, srts) = unzip [ srtTopRhs b r | (b,r) <- bs ]
+ bndrs = map fst bs
+ srts' = map (applyEnvList env) srts
+
+-- Shorting out indirections in SRTs: if a binding has an SRT with a single
+-- element in it, we just inline it with that element everywhere it occurs
+-- in other SRTs.
+--
+-- This is in a way a generalisation of the CafInfo. CafInfo says
+-- whether a top-level binding has *zero* CAF references, allowing us
+-- to omit it from SRTs. Here, we pick up bindings with *one* CAF
+-- reference, and inline its SRT everywhere it occurs. We could pass
+-- this information across module boundaries too, but we currently
+-- don't.
+
+maybeExtendEnv env bndr (StgRhsClosure _ _ _ ReEntrant (SRTEntries cafs) _ _)
+ | [one] <- varSetElems cafs
+ = extendVarEnv env bndr (applyEnv env one)
+maybeExtendEnv env bndr _ = env
+
+applyEnvList :: IdEnv Id -> [Id] -> [Id]
+applyEnvList env = map (applyEnv env)
+
+applyEnv env id = lookupVarEnv env id `orElse` id
+
+-- ---- Top-level right hand sides:
+
+srtTopRhs :: Id -> StgRhs -> (StgRhs, [Id])
+
+srtTopRhs binder rhs@(StgRhsCon _ _ _) = (rhs, [])
+srtTopRhs binder rhs@(StgRhsClosure _ _ _ _ (SRTEntries cafs) _ _)
+ = (srtRhs table rhs, elems)
+ where
+ elems = varSetElems cafs
+ table = mkVarEnv (zip elems [0..])