From 4d0dba944c3e103ddd59b7ab3df9ccd437985863 Mon Sep 17 00:00:00 2001
From: bjorn3 <bjorn3@users.noreply.github.com>
Date: Wed, 12 Dec 2018 16:01:34 +0100
Subject: [PATCH] Some refactorings for codegen_mono_items

---
 src/allocator.rs | 17 ++++++++++++++-
 src/common.rs    |  1 -
 src/lib.rs       | 54 +++++++++++++++++++++---------------------------
 3 files changed, 40 insertions(+), 32 deletions(-)

diff --git a/src/allocator.rs b/src/allocator.rs
index 8f2e7ca9569..9c444d8cfc5 100644
--- a/src/allocator.rs
+++ b/src/allocator.rs
@@ -13,7 +13,22 @@ use crate::prelude::*;
 use rustc::middle::allocator::AllocatorKind;
 use rustc_allocator::{AllocatorTy, ALLOCATOR_METHODS};
 
-pub fn codegen(module: &mut Module<impl Backend + 'static>, kind: AllocatorKind) {
+pub fn codegen(sess: &Session, module: &mut Module<impl Backend + 'static>) {
+    let any_dynamic_crate = sess
+        .dependency_formats
+        .borrow()
+        .iter()
+        .any(|(_, list)| {
+            use rustc::middle::dependency_format::Linkage;
+            list.iter().any(|&linkage| linkage == Linkage::Dynamic)
+        });
+    if any_dynamic_crate {
+    } else if let Some(kind) = *sess.allocator_kind.get() {
+        codegen_inner(module, kind);
+    }
+}
+
+pub fn codegen_inner(module: &mut Module<impl Backend + 'static>, kind: AllocatorKind) {
     let usize_ty = module.target_config().pointer_type();
 
     for method in ALLOCATOR_METHODS {
diff --git a/src/common.rs b/src/common.rs
index 8a6a40dff2a..1890caaf830 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -1,6 +1,5 @@
 use std::fmt;
 
-use rustc_codegen_ssa::traits::BackendTypes;
 use rustc_target::spec::{HasTargetSpec, Target};
 
 use cranelift_module::Module;
diff --git a/src/lib.rs b/src/lib.rs
index a2ebf470b3d..efedfbd1da3 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -80,6 +80,7 @@ mod prelude {
 
     pub use rustc_codegen_ssa::mir::operand::{OperandRef, OperandValue};
     pub use rustc_codegen_ssa::{CodegenResults, CompiledModule, ModuleKind};
+    pub use rustc_codegen_ssa::traits::*;
 
     pub use cranelift::codegen::ir::{
         condcodes::IntCC, function::Function, ExternalName, FuncRef, Inst, StackSlot,
@@ -200,6 +201,8 @@ impl CodegenBackend for CraneliftCodegenBackend {
                 .unwrap();
 
             codegen_mono_items(tcx, &mut jit_module, &mut log);
+            crate::allocator::codegen(tcx.sess, &mut jit_module);
+            jit_module.finalize_definitions();
 
             tcx.sess.abort_if_errors();
             println!("Compiled everything");
@@ -216,10 +219,9 @@ impl CodegenBackend for CraneliftCodegenBackend {
             jit_module.finish();
             ::std::process::exit(0);
         } else {
-            let isa = build_isa(tcx.sess);
             let mut faerie_module: Module<FaerieBackend> = Module::new(
                 FaerieBuilder::new(
-                    isa,
+                    build_isa(tcx.sess),
                     "some_file.o".to_string(),
                     FaerieTrapCollection::Disabled,
                     FaerieBuilder::default_libcall_names(),
@@ -232,6 +234,8 @@ impl CodegenBackend for CraneliftCodegenBackend {
             );
 
             codegen_mono_items(tcx, &mut faerie_module, &mut log);
+            crate::allocator::codegen(tcx.sess, &mut faerie_module);
+            faerie_module.finalize_definitions();
 
             tcx.sess.abort_if_errors();
 
@@ -339,36 +343,26 @@ fn codegen_mono_items<'a, 'tcx: 'a>(
         .flatten()
         .collect::<FxHashSet<(_, _)>>();
 
+    time("codegen mono items", move || {
+        for (&mono_item, &(_linkage, _vis)) in mono_items {
+            unimpl::try_unimpl(tcx, log, || {
+                base::trans_mono_item(tcx, module, &mut caches, &mut ccx, mono_item);
+            });
+        }
+
+        crate::main_shim::maybe_create_entry_wrapper(tcx, module);
+
+        ccx.finalize(tcx, module);
+    });
+}
+
+fn time<R>(name: &str, f: impl FnOnce() -> R) -> R {
+    println!("[{}] start", name);
     let before = ::std::time::Instant::now();
-    println!("[codegen mono items] start");
-
-    for (&mono_item, &(_linkage, _vis)) in mono_items {
-        unimpl::try_unimpl(tcx, log, || {
-            base::trans_mono_item(tcx, module, &mut caches, &mut ccx, mono_item);
-        });
-    }
-
-    crate::main_shim::maybe_create_entry_wrapper(tcx, module);
-
-    let any_dynamic_crate = tcx
-        .sess
-        .dependency_formats
-        .borrow()
-        .iter()
-        .any(|(_, list)| {
-            use rustc::middle::dependency_format::Linkage;
-            list.iter().any(|&linkage| linkage == Linkage::Dynamic)
-        });
-    if any_dynamic_crate {
-    } else if let Some(kind) = *tcx.sess.allocator_kind.get() {
-        allocator::codegen(module, kind);
-    }
-
-    ccx.finalize(tcx, module);
-    module.finalize_definitions();
-
+    let res = f();
     let after = ::std::time::Instant::now();
-    println!("[codegen mono items] end time: {:?}", after - before);
+    println!("[{}] end time: {:?}", name, after - before);
+    res
 }
 
 fn save_incremental<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>) {