lower move(expr) in plain closures

Author: TaKO8Ki

compiler/rustc_ast_lowering/src/expr.rs

@@ -2,6 +2,7 @@ use std::mem;
 use std::ops::ControlFlow;
 use std::sync::Arc;
 
+use rustc_ast::node_id::NodeMap;
 use rustc_ast::*;
 use rustc_ast_pretty::pprust::expr_to_string;
 use rustc_data_structures::stack::ensure_sufficient_stack;
@@ -29,6 +30,41 @@ use super::{
 use crate::errors::{InvalidLegacyConstGenericArg, UseConstGenericArg, YieldInClosure};
 use crate::{AllowReturnTypeNotation, FnDeclKind, ImplTraitPosition, TryBlockScope};
 
+struct MoveExprOccurrence<'a> {
+    id: NodeId,
+    move_kw_span: Span,
+    expr: &'a Expr,
+}
+
+struct MoveExprCollector<'a> {
+    occurrences: Vec<MoveExprOccurrence<'a>>,
+}
+
+impl<'a> MoveExprCollector<'a> {
+    fn collect(expr: &'a Expr) -> Vec<MoveExprOccurrence<'a>> {
+        let mut this = Self { occurrences: Vec::new() };
+        this.visit_expr(expr);
+        this.occurrences
+    }
+}
+
+impl<'a> Visitor<'a> for MoveExprCollector<'a> {
+    fn visit_expr(&mut self, expr: &'a Expr) {
+        match &expr.kind {
+            ExprKind::Move(inner, move_kw_span) => {
+                self.visit_expr(inner);
+                self.occurrences.push(MoveExprOccurrence {
+                    id: expr.id,
+                    move_kw_span: *move_kw_span,
+                    expr: inner,
+                });
+            }
+            ExprKind::Closure(..) | ExprKind::Gen(..) | ExprKind::ConstBlock(..) => {}
+            _ => walk_expr(self, expr),
+        }
+    }
+}
+
 struct WillCreateDefIdsVisitor {}
 
 impl<'v> rustc_ast::visit::Visitor<'v> for WillCreateDefIdsVisitor {
@@ -95,11 +131,12 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
                 ExprKind::ForLoop { pat, iter, body, label, kind } => {
                     return self.lower_expr_for(e, pat, iter, body, *label, *kind);
                 }
+                ExprKind::Closure(box closure) => return self.lower_expr_closure_expr(e, closure),
                 _ => (),
             }
 
             let expr_hir_id = self.lower_node_id(e.id);
-            let attrs = self.lower_attrs(expr_hir_id, &e.attrs, e.span, Target::from_expr(e));
+            self.lower_attrs(expr_hir_id, &e.attrs, e.span, Target::from_expr(e));
 
             let kind = match &e.kind {
                 ExprKind::Array(exprs) => hir::ExprKind::Array(self.lower_exprs(exprs)),
@@ -219,48 +256,34 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
                             self.dcx().span_delayed_bug(*move_kw_span, "invalid move(expr)"),
                         );
                     }
-                    self.dcx().emit_err(MoveExprOnlyInPlainClosures { span: *move_kw_span });
-                    hir::ExprKind::Err(
-                        self.dcx().span_delayed_bug(*move_kw_span, "invalid move(expr)"),
-                    )
+                    if let Some((ident, binding)) = self
+                        .move_expr_bindings
+                        .last()
+                        .and_then(|bindings| bindings.get(&e.id).copied())
+                    {
+                        hir::ExprKind::Path(hir::QPath::Resolved(
+                            None,
+                            self.arena.alloc(hir::Path {
+                                span: self.lower_span(e.span),
+                                res: Res::Local(binding),
+                                segments: arena_vec![
+                                    self;
+                                    hir::PathSegment::new(
+                                        self.lower_ident(ident),
+                                        self.next_id(),
+                                        Res::Local(binding),
+                                    )
+                                ],
+                            }),
+                        ))
+                    } else {
+                        self.dcx().emit_err(MoveExprOnlyInPlainClosures { span: *move_kw_span });
+                        hir::ExprKind::Err(
+                            self.dcx().span_delayed_bug(*move_kw_span, "invalid move(expr)"),
+                        )
+                    }
                 }
                 ExprKind::Use(expr, use_kw_span) => self.lower_expr_use(*use_kw_span, expr),
-                ExprKind::Closure(box Closure {
-                    binder,
-                    capture_clause,
-                    constness,
-                    coroutine_kind,
-                    movability,
-                    fn_decl,
-                    body,
-                    fn_decl_span,
-                    fn_arg_span,
-                }) => match coroutine_kind {
-                    Some(coroutine_kind) => self.lower_expr_coroutine_closure(
-                        binder,
-                        *capture_clause,
-                        e.id,
-                        expr_hir_id,
-                        *coroutine_kind,
-                        *constness,
-                        fn_decl,
-                        body,
-                        *fn_decl_span,
-                        *fn_arg_span,
-                    ),
-                    None => self.lower_expr_closure(
-                        attrs,
-                        binder,
-                        *capture_clause,
-                        e.id,
-                        *constness,
-                        *movability,
-                        fn_decl,
-                        body,
-                        *fn_decl_span,
-                        *fn_arg_span,
-                    ),
-                },
                 ExprKind::Gen(capture_clause, block, genblock_kind, decl_span) => {
                     let desugaring_kind = match genblock_kind {
                         GenBlockKind::Async => hir::CoroutineDesugaring::Async,
@@ -395,7 +418,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
 
                 ExprKind::Try(sub_expr) => self.lower_expr_try(e.span, sub_expr),
 
-                ExprKind::Paren(_) | ExprKind::ForLoop { .. } => {
+                ExprKind::Paren(_) | ExprKind::ForLoop { .. } | ExprKind::Closure(..) => {
                     unreachable!("already handled")
                 }
 
@@ -804,6 +827,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
             fn_arg_span: None,
             kind: hir::ClosureKind::Coroutine(coroutine_kind),
             constness: hir::Constness::NotConst,
+            explicit_captures: &[],
         }))
     }
 
@@ -1067,6 +1091,134 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         hir::ExprKind::Use(self.lower_expr(expr), self.lower_span(use_kw_span))
     }
 
+    fn lower_expr_closure_expr(&mut self, e: &Expr, closure: &Closure) -> hir::Expr<'hir> {
+        let expr_hir_id = self.lower_node_id(e.id);
+        let attrs = self.lower_attrs(expr_hir_id, &e.attrs, e.span, Target::from_expr(e));
+
+        match closure.coroutine_kind {
+            Some(coroutine_kind) => hir::Expr {
+                hir_id: expr_hir_id,
+                kind: self.lower_expr_coroutine_closure(
+                    &closure.binder,
+                    closure.capture_clause,
+                    e.id,
+                    expr_hir_id,
+                    coroutine_kind,
+                    closure.constness,
+                    &closure.fn_decl,
+                    &closure.body,
+                    closure.fn_decl_span,
+                    closure.fn_arg_span,
+                ),
+                span: self.lower_span(e.span),
+            },
+            None => self.lower_expr_plain_closure_with_move_exprs(
+                expr_hir_id,
+                attrs,
+                &closure.binder,
+                closure.capture_clause,
+                e.id,
+                closure.constness,
+                closure.movability,
+                &closure.fn_decl,
+                &closure.body,
+                closure.fn_decl_span,
+                closure.fn_arg_span,
+                e.span,
+            ),
+        }
+    }
+
+    fn lower_expr_plain_closure_with_move_exprs(
+        &mut self,
+        expr_hir_id: HirId,
+        attrs: &[rustc_hir::Attribute],
+        binder: &ClosureBinder,
+        capture_clause: CaptureBy,
+        closure_id: NodeId,
+        constness: Const,
+        movability: Movability,
+        decl: &FnDecl,
+        body: &Expr,
+        fn_decl_span: Span,
+        fn_arg_span: Span,
+        whole_span: Span,
+    ) -> hir::Expr<'hir> {
+        let occurrences = MoveExprCollector::collect(body);
+        if occurrences.is_empty() {
+            return hir::Expr {
+                hir_id: expr_hir_id,
+                kind: self.lower_expr_closure(
+                    attrs,
+                    binder,
+                    capture_clause,
+                    closure_id,
+                    constness,
+                    movability,
+                    decl,
+                    body,
+                    fn_decl_span,
+                    fn_arg_span,
+                    &[],
+                ),
+                span: self.lower_span(whole_span),
+            };
+        }
+
+        let mut bindings = NodeMap::default();
+        let mut lowered_occurrences = Vec::with_capacity(occurrences.len());
+        for (index, occurrence) in occurrences.iter().enumerate() {
+            let ident =
+                Ident::from_str_and_span(&format!("__move_expr_{index}"), occurrence.move_kw_span);
+            let (pat, binding) = self.pat_ident(occurrence.expr.span, ident);
+            bindings.insert(occurrence.id, (ident, binding));
+            lowered_occurrences.push((occurrence, pat, binding));
+        }
+
+        self.move_expr_bindings.push(bindings);
+        let mut stmts = Vec::with_capacity(lowered_occurrences.len());
+        for (occurrence, pat, _) in &lowered_occurrences {
+            let init = self.lower_expr(occurrence.expr);
+            stmts.push(self.stmt_let_pat(
+                None,
+                occurrence.expr.span,
+                Some(init),
+                *pat,
+                hir::LocalSource::Normal,
+            ));
+        }
+
+        let explicit_captures = self.arena.alloc_from_iter(lowered_occurrences.iter().map(
+            |(occurrence, _, binding)| hir::ExplicitCapture {
+                var_hir_id: *binding,
+                origin_span: self.lower_span(occurrence.move_kw_span),
+            },
+        ));
+
+        let closure_expr = self.arena.alloc(hir::Expr {
+            hir_id: expr_hir_id,
+            kind: self.lower_expr_closure(
+                attrs,
+                binder,
+                capture_clause,
+                closure_id,
+                constness,
+                movability,
+                decl,
+                body,
+                fn_decl_span,
+                fn_arg_span,
+                explicit_captures,
+            ),
+            span: self.lower_span(whole_span),
+        });
+        self.move_expr_bindings.pop();
+
+        let stmts = self.arena.alloc_from_iter(stmts);
+        let block = self.block_all(whole_span, stmts, Some(closure_expr));
+        self.expr(whole_span, hir::ExprKind::Block(block, None))
+    }
+
     fn lower_expr_closure(
         &mut self,
         attrs: &[rustc_hir::Attribute],
@@ -1079,6 +1231,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
         body: &Expr,
         fn_decl_span: Span,
         fn_arg_span: Span,
+        explicit_captures: &'hir [hir::ExplicitCapture],
     ) -> hir::ExprKind<'hir> {
         let closure_def_id = self.local_def_id(closure_id);
         let (binder_clause, generic_params) = self.lower_closure_binder(binder);
@@ -1120,6 +1273,7 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
             fn_arg_span: Some(self.lower_span(fn_arg_span)),
             kind: closure_kind,
             constness: self.lower_constness(constness),
+            explicit_captures,
         });
 
         hir::ExprKind::Closure(c)
@@ -1241,7 +1395,8 @@ impl<'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'_, 'hir, R> {
             // knows that a `FnDecl` output type like `-> &str` actually means
             // "coroutine that returns &str", rather than directly returning a `&str`.
             kind: hir::ClosureKind::CoroutineClosure(coroutine_desugaring),
-            constness: self.lower_constness(constness),
+            constness: hir::Constness::NotConst,
+            explicit_captures: &[],
         });
         hir::ExprKind::Closure(c)
     }

compiler/rustc_ast_lowering/src/lib.rs

@@ -148,6 +148,7 @@ struct LoweringContext<'a, 'hir, R> {
     allow_async_fn_traits: Arc<[Symbol]>,
 
     delayed_lints: Vec<DelayedLint>,
+    move_expr_bindings: Vec<NodeMap<(Ident, HirId)>>,
 
     attribute_parser: AttributeParser<'hir>,
 }
@@ -205,6 +206,7 @@ impl<'a, 'hir, R: ResolverAstLoweringExt<'hir>> LoweringContext<'a, 'hir, R> {
             // interact with `gen`/`async gen` blocks
             allow_async_iterator: [sym::gen_future, sym::async_iterator].into(),
 
+            move_expr_bindings: Vec::new(),
             attribute_parser: AttributeParser::new(
                 tcx.sess,
                 tcx.features(),

compiler/rustc_hir/src/hir.rs

@@ -1690,6 +1690,13 @@ pub struct Closure<'hir> {
     /// The span of the argument block `|...|`
     pub fn_arg_span: Option<Span>,
     pub kind: ClosureKind,
+    pub explicit_captures: &'hir [ExplicitCapture],
+}
+
+#[derive(Debug, Clone, Copy, HashStable_Generic)]
+pub struct ExplicitCapture {
+    pub var_hir_id: HirId,
+    pub origin_span: Span,
 }
 
 #[derive(Clone, PartialEq, Eq, Debug, Copy, Hash, HashStable_Generic, Encodable, Decodable)]

compiler/rustc_hir/src/intravisit.rs

@@ -894,6 +894,7 @@ pub fn walk_expr<'v, V: Visitor<'v>>(visitor: &mut V, expression: &'v Expr<'v>)
             fn_arg_span: _,
             kind: _,
             constness: _,
+            explicit_captures: _,
         }) => {
             walk_list!(visitor, visit_generic_param, bound_generic_params);
             try_visit!(visitor.visit_fn(FnKind::Closure, fn_decl, body, *span, def_id));

compiler/rustc_hir_pretty/src/lib.rs

@@ -1643,6 +1643,7 @@ impl<'a> State<'a> {
                 fn_arg_span: _,
                 kind: _,
                 def_id: _,
+                explicit_captures: _,
             }) => {
                 self.print_closure_binder(binder, bound_generic_params);
                 self.print_constness(constness);

tests/ui/move-expr/outside-plain-closure.rs

@@ -0,0 +1,7 @@
+#![allow(incomplete_features)]
+#![feature(move_expr)]
+
+fn main() {
+    let _ = move(String::from("nope"));
+    //~^ ERROR `move(expr)` is only supported in plain closures
+}

tests/ui/move-expr/outside-plain-closure.stderr

@@ -0,0 +1,8 @@
+error: `move(expr)` is only supported in plain closures
+  --> $DIR/outside-plain-closure.rs:5:13
+   |
+LL |     let _ = move(String::from("nope"));
+   |             ^^^^
+
+error: aborting due to 1 previous error
+

tests/ui/move-expr/plain-closure.rs

@@ -0,0 +1,12 @@
+//@ check-pass
+#![allow(incomplete_features)]
+#![feature(move_expr)]
+
+fn main() {
+    let s = String::from("hello");
+    let c = || {
+        let t = move(s);
+        println!("{}", t.len());
+    };
+    c();
+}