From 5c9a8b504146cc9f755c4e7d73b3693cb781021f Mon Sep 17 00:00:00 2001
From: varkor <github@varkor.com>
Date: Fri, 19 Jan 2018 00:32:58 +0000
Subject: [PATCH] Add instance_def_size_estimate query

---
 src/librustc/dep_graph/dep_node.rs |  1 +
 src/librustc/mir/mono.rs           |  3 +--
 src/librustc/ty/maps/config.rs     |  6 ++++++
 src/librustc/ty/maps/mod.rs        | 10 ++++++++++
 src/librustc/ty/maps/plumbing.rs   |  1 +
 src/librustc/ty/mod.rs             | 14 ++++++++++++++
 6 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/src/librustc/dep_graph/dep_node.rs b/src/librustc/dep_graph/dep_node.rs
index 14f54fbffac2..1de9091b5df7 100644
--- a/src/librustc/dep_graph/dep_node.rs
+++ b/src/librustc/dep_graph/dep_node.rs
@@ -638,6 +638,7 @@ define_dep_nodes!( <'tcx>
     [input] TargetFeaturesWhitelist,
     [] TargetFeaturesEnabled(DefId),
 
+    [] InstanceDefSizeEstimate { instance_def: InstanceDef<'tcx> },
 );
 
 trait DepNodeParams<'a, 'gcx: 'tcx + 'a, 'tcx: 'a> : fmt::Debug {
diff --git a/src/librustc/mir/mono.rs b/src/librustc/mir/mono.rs
index bbef045a305a..2af2219d2678 100644
--- a/src/librustc/mir/mono.rs
+++ b/src/librustc/mir/mono.rs
@@ -31,8 +31,7 @@ impl<'tcx> MonoItem<'tcx> {
             MonoItem::Fn(instance) => {
                 // Estimate the size of a function based on how many statements
                 // it contains.
-                let mir = tcx.instance_mir(instance.def);
-                mir.basic_blocks().iter().map(|bb| bb.statements.len()).sum()
+                tcx.instance_def_size_estimate(instance.def)
             },
             // Conservatively estimate the size of a static declaration
             // or assembly to be 1.
diff --git a/src/librustc/ty/maps/config.rs b/src/librustc/ty/maps/config.rs
index 8dedcb24c2fb..b0adfe208796 100644
--- a/src/librustc/ty/maps/config.rs
+++ b/src/librustc/ty/maps/config.rs
@@ -637,6 +637,12 @@ impl<'tcx> QueryDescription<'tcx> for queries::target_features_whitelist<'tcx> {
     }
 }
 
+impl<'tcx> QueryDescription<'tcx> for queries::instance_def_size_estimate<'tcx> {
+    fn describe(tcx: TyCtxt, def: ty::InstanceDef<'tcx>) -> String {
+        format!("estimating size for `{}`", tcx.item_path_str(def.def_id()))
+    }
+}
+
 macro_rules! impl_disk_cacheable_query(
     ($query_name:ident, |$key:tt| $cond:expr) => {
         impl<'tcx> QueryDescription<'tcx> for queries::$query_name<'tcx> {
diff --git a/src/librustc/ty/maps/mod.rs b/src/librustc/ty/maps/mod.rs
index e7e92b8a4288..6c79f6a62fa0 100644
--- a/src/librustc/ty/maps/mod.rs
+++ b/src/librustc/ty/maps/mod.rs
@@ -365,6 +365,9 @@ define_maps! { <'tcx>
         target_features_whitelist_node(CrateNum) -> Rc<FxHashSet<String>>,
     [] fn target_features_enabled: TargetFeaturesEnabled(DefId) -> Rc<Vec<String>>,
 
+    // Get an estimate of the size of an InstanceDef based on its MIR for CGU partitioning.
+    [] fn instance_def_size_estimate: instance_def_size_estimate_dep_node(ty::InstanceDef<'tcx>)
+        -> usize,
 }
 
 //////////////////////////////////////////////////////////////////////
@@ -514,3 +517,10 @@ fn substitute_normalize_and_test_predicates_node<'tcx>(key: (DefId, &'tcx Substs
 fn target_features_whitelist_node<'tcx>(_: CrateNum) -> DepConstructor<'tcx> {
     DepConstructor::TargetFeaturesWhitelist
 }
+
+fn instance_def_size_estimate_dep_node<'tcx>(instance_def: ty::InstanceDef<'tcx>)
+                                              -> DepConstructor<'tcx> {
+    DepConstructor::InstanceDefSizeEstimate {
+        instance_def
+    }
+}
diff --git a/src/librustc/ty/maps/plumbing.rs b/src/librustc/ty/maps/plumbing.rs
index d670ecc2691a..f35693c702b3 100644
--- a/src/librustc/ty/maps/plumbing.rs
+++ b/src/librustc/ty/maps/plumbing.rs
@@ -761,6 +761,7 @@ pub fn force_from_dep_node<'a, 'gcx, 'lcx>(tcx: TyCtxt<'a, 'gcx, 'lcx>,
         DepKind::EraseRegionsTy |
         DepKind::NormalizeTy |
         DepKind::SubstituteNormalizeAndTestPredicates |
+        DepKind::InstanceDefSizeEstimate |
 
         // This one should never occur in this context
         DepKind::Null => {
diff --git a/src/librustc/ty/mod.rs b/src/librustc/ty/mod.rs
index 2b4d2c80c6f9..1d64e7bae913 100644
--- a/src/librustc/ty/mod.rs
+++ b/src/librustc/ty/mod.rs
@@ -2669,6 +2669,19 @@ fn crate_hash<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
     tcx.hir.crate_hash
 }
 
+fn instance_def_size_estimate<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
+                                        instance_def: InstanceDef<'tcx>)
+                                        -> usize {
+    match instance_def {
+        InstanceDef::Item(def_id) => {
+            let mir = tcx.optimized_mir(def_id);
+            mir.basic_blocks().iter().map(|bb| bb.statements.len()).sum()
+        },
+        // Estimate the size of compiler-generated shims to be 1.
+        _ => 1
+    }
+}
+
 pub fn provide(providers: &mut ty::maps::Providers) {
     context::provide(providers);
     erase_regions::provide(providers);
@@ -2686,6 +2699,7 @@ pub fn provide(providers: &mut ty::maps::Providers) {
         original_crate_name,
         crate_hash,
         trait_impls_of: trait_def::trait_impls_of_provider,
+        instance_def_size_estimate,
         ..*providers
     };
 }