From acbfa06ea2aa298de1591c8bb9a483151f254c2d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=96mer=20Sinan=20A=C4=9Facan?= <omeragacan@gmail.com>
Date: Fri, 29 May 2020 12:06:29 +0300
Subject: [PATCH] Mark blocks that call cold funs as cold (#1021)

---
 src/abi/mod.rs | 13 +++++++++++--
 src/base.rs    |  1 +
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/abi/mod.rs b/src/abi/mod.rs
index f4a2fd566aca..8f59fa91bfb4 100644
--- a/src/abi/mod.rs
+++ b/src/abi/mod.rs
@@ -4,6 +4,7 @@ mod pass_mode;
 mod returning;
 
 use rustc_target::spec::abi::Abi;
+use rustc_middle::middle::codegen_fn_attrs::CodegenFnAttrFlags;
 
 use cranelift_codegen::ir::AbiParam;
 
@@ -431,6 +432,7 @@ pub(crate) fn codegen_fn_prelude<'tcx>(
 pub(crate) fn codegen_terminator_call<'tcx>(
     fx: &mut FunctionCx<'_, 'tcx, impl Backend>,
     span: Span,
+    current_block: Block,
     func: &Operand<'tcx>,
     args: &[Operand<'tcx>],
     destination: Option<(Place<'tcx>, BasicBlock)>,
@@ -440,8 +442,6 @@ pub(crate) fn codegen_terminator_call<'tcx>(
         .tcx
         .normalize_erasing_late_bound_regions(ParamEnv::reveal_all(), &fn_ty.fn_sig(fx.tcx));
 
-    // FIXME mark the current block as cold when calling a `#[cold]` function.
-
     let destination = destination.map(|(place, bb)| (trans_place(fx, place), bb));
 
     // Handle special calls like instrinsics and empty drop glue.
@@ -479,6 +479,15 @@ pub(crate) fn codegen_terminator_call<'tcx>(
         None
     };
 
+    let is_cold =
+        instance.map(|inst|
+            fx.tcx.codegen_fn_attrs(inst.def_id())
+                .flags.contains(CodegenFnAttrFlags::COLD))
+                .unwrap_or(false);
+    if is_cold {
+        fx.cold_blocks.insert(current_block);
+    }
+
     // Unpack arguments tuple for closures
     let args = if fn_sig.abi == Abi::RustCall {
         assert_eq!(args.len(), 2, "rust-call abi requires two arguments");
diff --git a/src/base.rs b/src/base.rs
index b3751a020294..201b0f7a181a 100644
--- a/src/base.rs
+++ b/src/base.rs
@@ -300,6 +300,7 @@ fn codegen_fn_content(fx: &mut FunctionCx<'_, '_, impl Backend>) {
                 fx.tcx.sess.time("codegen call", || crate::abi::codegen_terminator_call(
                     fx,
                     bb_data.terminator().source_info.span,
+                    block,
                     func,
                     args,
                     *destination,