@@ -91,6 +91,10 @@ pub struct FunctionCx<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>> {
9191 /// Cached terminate upon unwinding block and its reason
9292 terminate_block : Option < ( Bx :: BasicBlock , UnwindTerminateReason ) > ,
9393
94+ /// A bool flag for each basic block indicating whether it is a cold block.
95+ /// A cold block is a block that is unlikely to be executed at runtime.
96+ cold_blocks : IndexVec < mir:: BasicBlock , bool > ,
97+
9498 /// The location where each MIR arg/var/tmp/ret is stored. This is
9599 /// usually an `PlaceRef` representing an alloca, but not always:
96100 /// sometimes we can skip the alloca and just store the value
@@ -207,6 +211,7 @@ pub fn codegen_mir<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
207211 cleanup_kinds,
208212 landing_pads : IndexVec :: from_elem ( None , & mir. basic_blocks ) ,
209213 funclets : IndexVec :: from_fn_n ( |_| None , mir. basic_blocks . len ( ) ) ,
214+ cold_blocks : find_cold_blocks ( cx. tcx ( ) , mir) ,
210215 locals : locals:: Locals :: empty ( ) ,
211216 debug_context,
212217 per_local_var_debug_info : None ,
@@ -477,3 +482,39 @@ fn arg_local_refs<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
477482
478483 args
479484}
485+
486+ fn find_cold_blocks < ' tcx > (
487+ tcx : TyCtxt < ' tcx > ,
488+ mir : & mir:: Body < ' tcx > ,
489+ ) -> IndexVec < mir:: BasicBlock , bool > {
490+ let local_decls = & mir. local_decls ;
491+
492+ let mut cold_blocks: IndexVec < mir:: BasicBlock , bool > =
493+ IndexVec :: from_elem ( false , & mir. basic_blocks ) ;
494+
495+ // Traverse all basic blocks from end of the function to the start.
496+ for ( bb, bb_data) in traversal:: postorder ( mir) {
497+ let terminator = bb_data. terminator ( ) ;
498+
499+ // If a BB ends with a call to a cold function, mark it as cold.
500+ if let mir:: TerminatorKind :: Call { ref func, .. } = terminator. kind
501+ && let ty:: FnDef ( def_id, ..) = * func. ty ( local_decls, tcx) . kind ( )
502+ && let attrs = tcx. codegen_fn_attrs ( def_id)
503+ && attrs. flags . contains ( CodegenFnAttrFlags :: COLD )
504+ {
505+ cold_blocks[ bb] = true ;
506+ continue ;
507+ }
508+
509+ // If all successors of a BB are cold and there's at least one of them, mark this BB as cold
510+ let mut succ = terminator. successors ( ) ;
511+ if let Some ( first) = succ. next ( )
512+ && cold_blocks[ first]
513+ && succ. all ( |s| cold_blocks[ s] )
514+ {
515+ cold_blocks[ bb] = true ;
516+ }
517+ }
518+
519+ cold_blocks
520+ }
0 commit comments