Differentiate between closure and function bodies

Author: oli-obk

src/librustc/hir/map/mod.rs

@@ -455,11 +455,20 @@ impl<'hir> Map<'hir> {
             Node::AnonConst(_) => {
                 BodyOwnerKind::Const
             }
+            Node::Variant(&Spanned { node: VariantKind { data: VariantData::Tuple(..), .. }, .. }) |
+            Node::StructCtor(..) |
+            Node::Item(&Item { node: ItemKind::Fn(..), .. }) |
+            Node::TraitItem(&TraitItem { node: TraitItemKind::Method(..), .. }) |
+            Node::ImplItem(&ImplItem { node: ImplItemKind::Method(..), .. }) => {
+                BodyOwnerKind::Fn
+            }
             Node::Item(&Item { node: ItemKind::Static(_, m, _), .. }) => {
                 BodyOwnerKind::Static(m)
             }
-            // Default to function if it's not a constant or static.
-            _ => BodyOwnerKind::Fn
+            Node::Expr(&Expr { node: ExprKind::Closure(..), .. }) => {
+                BodyOwnerKind::Closure
+            }
+            node => bug!("{:#?} is not a body node", node),
         }
     }
 

src/librustc/hir/mod.rs

@@ -1270,13 +1270,25 @@ pub enum BodyOwnerKind {
     /// Functions and methods.
     Fn,
 
+    /// Closures
+    Closure,
+
     /// Constants and associated constants.
     Const,
 
     /// Initializer of a `static` item.
     Static(Mutability),
 }
 
+impl BodyOwnerKind {
+    pub fn is_fn_or_closure(self) -> bool {
+        match self {
+            BodyOwnerKind::Fn | BodyOwnerKind::Closure => true,
+            BodyOwnerKind::Const | BodyOwnerKind::Static(_) => false,
+        }
+    }
+}
+
 /// A constant (expression) that's not an item or associated item,
 /// but needs its own `DefId` for type-checking, const-eval, etc.
 /// These are usually found nested inside types (e.g., array lengths)

src/librustc/middle/region.rs

@@ -1268,8 +1268,8 @@ impl<'a, 'tcx> Visitor<'tcx> for RegionResolutionVisitor<'a, 'tcx> {
 
         // The body of the every fn is a root scope.
         self.cx.parent = self.cx.var_parent;
-        if let hir::BodyOwnerKind::Fn = self.tcx.hir().body_owner_kind(owner_id) {
-            self.visit_expr(&body.value);
+        if self.tcx.hir().body_owner_kind(owner_id).is_fn_or_closure() {
+            self.visit_expr(&body.value)
         } else {
             // Only functions have an outer terminating (drop) scope, while
             // temporaries in constant initializers may be 'static, but only

src/librustc_mir/borrow_check/mod.rs

@@ -163,10 +163,7 @@ fn do_mir_borrowck<'a, 'gcx, 'tcx>(
         |bd, i| DebugFormatted::new(&bd.move_data().move_paths[i]),
     ));
 
-    let locals_are_invalidated_at_exit = match tcx.hir().body_owner_kind(id) {
-            hir::BodyOwnerKind::Const | hir::BodyOwnerKind::Static(_) => false,
-            hir::BodyOwnerKind::Fn => true,
-    };
+    let locals_are_invalidated_at_exit = tcx.hir().body_owner_kind(id).is_fn_or_closure();
     let borrow_set = Rc::new(BorrowSet::build(
             tcx, mir, locals_are_invalidated_at_exit, &mdpe.move_data));
 

src/librustc_mir/borrow_check/nll/universal_regions.rs

@@ -476,6 +476,7 @@ impl<'cx, 'gcx, 'tcx> UniversalRegionsBuilder<'cx, 'gcx, 'tcx> {
         let closure_base_def_id = tcx.closure_base_def_id(self.mir_def_id);
 
         match tcx.hir().body_owner_kind(self.mir_node_id) {
+            BodyOwnerKind::Closure |
             BodyOwnerKind::Fn => {
                 let defining_ty = if self.mir_def_id == closure_base_def_id {
                     tcx.type_of(closure_base_def_id)

src/librustc_mir/build/mod.rs

@@ -75,7 +75,7 @@ pub fn mir_build<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId) -> Mir<'t
         let cx = Cx::new(&infcx, id);
         let mut mir = if cx.tables().tainted_by_errors {
             build::construct_error(cx, body_id)
-        } else if let hir::BodyOwnerKind::Fn = cx.body_owner_kind {
+        } else if cx.body_owner_kind.is_fn_or_closure() {
             // fetch the fully liberated fn signature (that is, all bound
             // types/lifetimes replaced)
             let fn_hir_id = tcx.hir().node_to_hir_id(id);

src/librustc_mir/build/scope.rs

@@ -613,6 +613,7 @@ impl<'a, 'gcx, 'tcx> Builder<'a, 'gcx, 'tcx> {
             hir::BodyOwnerKind::Static(_) =>
                 // No need to free storage in this context.
                 None,
+            hir::BodyOwnerKind::Closure |
             hir::BodyOwnerKind::Fn =>
                 Some(self.topmost_scope()),
         }

src/librustc_mir/hair/cx/mod.rs

@@ -61,6 +61,7 @@ impl<'a, 'gcx, 'tcx> Cx<'a, 'gcx, 'tcx> {
         let constness = match body_owner_kind {
             hir::BodyOwnerKind::Const |
             hir::BodyOwnerKind::Static(_) => hir::Constness::Const,
+            hir::BodyOwnerKind::Closure |
             hir::BodyOwnerKind::Fn => hir::Constness::NotConst,
         };
 

src/librustc_mir/transform/inline.rs

@@ -1,6 +1,5 @@
 //! Inlining pass for MIR functions
 
-use rustc::hir;
 use rustc::hir::CodegenFnAttrFlags;
 use rustc::hir::def_id::DefId;
 
@@ -74,15 +73,12 @@ impl<'a, 'tcx> Inliner<'a, 'tcx> {
 
         // Only do inlining into fn bodies.
         let id = self.tcx.hir().as_local_node_id(self.source.def_id).unwrap();
-        let body_owner_kind = self.tcx.hir().body_owner_kind(id);
-
-        if let (hir::BodyOwnerKind::Fn, None) = (body_owner_kind, self.source.promoted) {
-
+        if self.tcx.hir().body_owner_kind(id).is_fn_or_closure() && self.source.promoted.is_none() {
             for (bb, bb_data) in caller_mir.basic_blocks().iter_enumerated() {
                 if let Some(callsite) = self.get_valid_function_call(bb,
-                                                                     bb_data,
-                                                                     caller_mir,
-                                                                     param_env) {
+                                                                    bb_data,
+                                                                    caller_mir,
+                                                                    param_env) {
                     callsites.push_back(callsite);
                 }
             }

src/librustc_mir/transform/qualify_consts.rs

@@ -1152,6 +1152,7 @@ impl MirPass for QualifyAndPromoteConstants {
         let id = tcx.hir().as_local_node_id(def_id).unwrap();
         let mut const_promoted_temps = None;
         let mode = match tcx.hir().body_owner_kind(id) {
+            hir::BodyOwnerKind::Closure => Mode::Fn,
             hir::BodyOwnerKind::Fn => {
                 if tcx.is_const_fn(def_id) {
                     Mode::ConstFn

src/librustc_mir/util/pretty.rs

@@ -573,6 +573,7 @@ fn write_mir_sig(tcx: TyCtxt, src: MirSource, mir: &Mir, w: &mut dyn Write) -> i
     let body_owner_kind = tcx.hir().body_owner_kind(id);
     match (body_owner_kind, src.promoted) {
         (_, Some(i)) => write!(w, "{:?} in", i)?,
+        (hir::BodyOwnerKind::Closure, _) |
         (hir::BodyOwnerKind::Fn, _) => write!(w, "fn")?,
         (hir::BodyOwnerKind::Const, _) => write!(w, "const")?,
         (hir::BodyOwnerKind::Static(hir::MutImmutable), _) => write!(w, "static")?,
@@ -585,6 +586,7 @@ fn write_mir_sig(tcx: TyCtxt, src: MirSource, mir: &Mir, w: &mut dyn Write) -> i
     })?;
 
     match (body_owner_kind, src.promoted) {
+        (hir::BodyOwnerKind::Closure, None) |
         (hir::BodyOwnerKind::Fn, None) => {
             write!(w, "(")?;
 

src/librustc_passes/rvalue_promotion.rs

@@ -191,6 +191,7 @@ impl<'a, 'tcx> CheckCrateVisitor<'a, 'tcx> {
         self.in_static = false;
 
         match self.tcx.hir().body_owner_kind(item_id) {
+            hir::BodyOwnerKind::Closure |
             hir::BodyOwnerKind::Fn => self.in_fn = true,
             hir::BodyOwnerKind::Static(_) => self.in_static = true,
             _ => {}