Suggest move in nested closure when appropriate

Fix rust-lang#64008.
matthiaskrgr · Jan 11, 2023 · 00f821f · 00f821f
1 parent b22c152
commit 00f821f
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 21 deletions.
diff --git a/compiler/rustc_borrowck/src/diagnostics/region_errors.rs b/compiler/rustc_borrowck/src/diagnostics/region_errors.rs
@@ -474,10 +474,12 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
         let err = FnMutError {
             span: *span,
             ty_err: match output_ty.kind() {
-                ty::Closure(_, _) => FnMutReturnTypeErr::ReturnClosure { span: *span },
                 ty::Generator(def, ..) if self.infcx.tcx.generator_is_async(*def) => {
                     FnMutReturnTypeErr::ReturnAsyncBlock { span: *span }
                 }
+                _ if output_ty.contains_closure() => {
+                    FnMutReturnTypeErr::ReturnClosure { span: *span }
+                }
                 _ => FnMutReturnTypeErr::ReturnRef { span: *span },
             },
         };
@@ -867,7 +869,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
     fn suggest_move_on_borrowing_closure(&self, diag: &mut Diagnostic) {
         let map = self.infcx.tcx.hir();
         let body_id = map.body_owned_by(self.mir_def_id());
-        let expr = &map.body(body_id).value;
+        let expr = &map.body(body_id).value.peel_blocks();
         let mut closure_span = None::<rustc_span::Span>;
         match expr.kind {
             hir::ExprKind::MethodCall(.., args, _) => {
@@ -882,20 +884,14 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
                     }
                 }
             }
-            hir::ExprKind::Block(blk, _) => {
-                if let Some(expr) = blk.expr {
-                    // only when the block is a closure
-                    if let hir::ExprKind::Closure(hir::Closure {
-                        capture_clause: hir::CaptureBy::Ref,
-                        body,
-                        ..
-                    }) = expr.kind
-                    {
-                        let body = map.body(*body);
-                        if !matches!(body.generator_kind, Some(hir::GeneratorKind::Async(..))) {
-                            closure_span = Some(expr.span.shrink_to_lo());
-                        }
-                    }
+            hir::ExprKind::Closure(hir::Closure {
+                capture_clause: hir::CaptureBy::Ref,
+                body,
+                ..
+            }) => {
+                let body = map.body(*body);
+                if !matches!(body.generator_kind, Some(hir::GeneratorKind::Async(..))) {
+                    closure_span = Some(expr.span.shrink_to_lo());
                 }
             }
             _ => {}

diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
@@ -1937,6 +1937,28 @@ impl<'tcx> Ty<'tcx> {
         cf.is_break()
     }
 
+    /// Checks whether a type recursively contains any closure
+    ///
+    /// Example: `Option<[[email protected]:4:20]>` returns true
+    pub fn contains_closure(self) -> bool {
+        struct ContainsClosureVisitor;
+
+        impl<'tcx> TypeVisitor<'tcx> for ContainsClosureVisitor {
+            type BreakTy = ();
+
+            fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<Self::BreakTy> {
+                if let ty::Closure(_, _) = t.kind() {
+                    ControlFlow::BREAK
+                } else {
+                    t.super_visit_with(self)
+                }
+            }
+        }
+
+        let cf = self.visit_with(&mut ContainsClosureVisitor);
+        cf.is_break()
+    }
+
     /// Returns the type and mutability of `*ty`.
     ///
     /// The parameter `explicit` indicates if this is an *explicit* dereference.

diff --git a/tests/ui/borrowck/issue-95079-missing-move-in-nested-closure.fixed b/tests/ui/borrowck/issue-95079-missing-move-in-nested-closure.fixed
@@ -0,0 +1,26 @@
+// run-rustfix
+#![allow(dead_code, path_statements)]
+fn foo1(s: &str) -> impl Iterator<Item = String> + '_ {
+    None.into_iter()
+        .flat_map(move |()| s.chars().map(move |c| format!("{}{}", c, s)))
+        //~^ ERROR captured variable cannot escape `FnMut` closure body
+        //~| HELP consider adding 'move' keyword before the nested closure
+}
+
+fn foo2(s: &str) -> impl Sized + '_ {
+    move |()| s.chars().map(move |c| format!("{}{}", c, s))
+    //~^ ERROR lifetime may not live long enough
+    //~| HELP consider adding 'move' keyword before the nested closure
+}
+
+pub struct X;
+pub fn foo3<'a>(
+    bar: &'a X,
+) -> impl Iterator<Item = ()> + 'a {
+    Some(()).iter().flat_map(move |()| {
+        Some(()).iter().map(move |()| { bar; }) //~ ERROR captured variable cannot escape
+        //~^ HELP consider adding 'move' keyword before the nested closure
+    })
+}
+
+fn main() {}
diff --git a/tests/ui/borrowck/issue-95079-missing-move-in-nested-closure.rs b/tests/ui/borrowck/issue-95079-missing-move-in-nested-closure.rs
@@ -1,3 +1,5 @@
+// run-rustfix
+#![allow(dead_code, path_statements)]
 fn foo1(s: &str) -> impl Iterator<Item = String> + '_ {
     None.into_iter()
         .flat_map(move |()| s.chars().map(|c| format!("{}{}", c, s)))
@@ -11,4 +13,14 @@ fn foo2(s: &str) -> impl Sized + '_ {
     //~| HELP consider adding 'move' keyword before the nested closure
 }
 
+pub struct X;
+pub fn foo3<'a>(
+    bar: &'a X,
+) -> impl Iterator<Item = ()> + 'a {
+    Some(()).iter().flat_map(move |()| {
+        Some(()).iter().map(|()| { bar; }) //~ ERROR captured variable cannot escape
+        //~^ HELP consider adding 'move' keyword before the nested closure
+    })
+}
+
 fn main() {}
diff --git a/tests/ui/borrowck/issue-95079-missing-move-in-nested-closure.stderr b/tests/ui/borrowck/issue-95079-missing-move-in-nested-closure.stderr
@@ -1,13 +1,13 @@
 error: captured variable cannot escape `FnMut` closure body
-  --> $DIR/issue-95079-missing-move-in-nested-closure.rs:3:29
+  --> $DIR/issue-95079-missing-move-in-nested-closure.rs:5:29
    |
 LL | fn foo1(s: &str) -> impl Iterator<Item = String> + '_ {
    |         - variable defined here
 LL |     None.into_iter()
 LL |         .flat_map(move |()| s.chars().map(|c| format!("{}{}", c, s)))
    |                           - -^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |                           | |
-   |                           | returns a reference to a captured variable which escapes the closure body
+   |                           | returns a closure that contains a reference to a captured variable, which then escapes the closure body
    |                           | variable captured here
    |                           inferred to be a `FnMut` closure
    |
@@ -19,12 +19,12 @@ LL |         .flat_map(move |()| s.chars().map(move |c| format!("{}{}", c, s)))
    |                                           ++++
 
 error: lifetime may not live long enough
-  --> $DIR/issue-95079-missing-move-in-nested-closure.rs:9:15
+  --> $DIR/issue-95079-missing-move-in-nested-closure.rs:11:15
    |
 LL |     move |()| s.chars().map(|c| format!("{}{}", c, s))
    |     --------- ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ returning this value requires that `'1` must outlive `'2`
    |     |       |
-   |     |       return type of closure `Map<Chars<'_>, [closure@$DIR/issue-95079-missing-move-in-nested-closure.rs:9:29: 9:32]>` contains a lifetime `'2`
+   |     |       return type of closure `Map<Chars<'_>, [closure@$DIR/issue-95079-missing-move-in-nested-closure.rs:11:29: 11:32]>` contains a lifetime `'2`
    |     lifetime `'1` represents this closure's body
    |
    = note: closure implements `Fn`, so references to captured variables can't escape the closure
@@ -33,5 +33,26 @@ help: consider adding 'move' keyword before the nested closure
 LL |     move |()| s.chars().map(move |c| format!("{}{}", c, s))
    |                             ++++
 
-error: aborting due to 2 previous errors
+error: captured variable cannot escape `FnMut` closure body
+  --> $DIR/issue-95079-missing-move-in-nested-closure.rs:21:9
+   |
+LL |     bar: &'a X,
+   |     --- variable defined here
+LL | ) -> impl Iterator<Item = ()> + 'a {
+LL |     Some(()).iter().flat_map(move |()| {
+   |                                      - inferred to be a `FnMut` closure
+LL |         Some(()).iter().map(|()| { bar; })
+   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^---^^^^
+   |         |                          |
+   |         |                          variable captured here
+   |         returns a closure that contains a reference to a captured variable, which then escapes the closure body
+   |
+   = note: `FnMut` closures only have access to their captured variables while they are executing...
+   = note: ...therefore, they cannot allow references to captured variables to escape
+help: consider adding 'move' keyword before the nested closure
+   |
+LL |         Some(()).iter().map(move |()| { bar; })
+   |                             ++++
+
+error: aborting due to 3 previous errors