compiler/rustc_ast_lowering/src/expr/closure.rs

use rustc_ast::*;
use rustc_attr_ir::find_attr;
use rustc_attr_ir::target::Target;
use rustc_hir as hir;
use rustc_hir::HirId;
use rustc_span::{Span, span_bug};

use super::{LoweringContext, MoveExprState};
use crate::FnDeclKind;
use crate::diagnostics::{ClosureCannotBeStatic, CoroutineTooManyParameters};

impl<'hir> LoweringContext<'_, 'hir> {
    // Entry point for `ExprKind::Closure`. Plain closures go through
    // `lower_expr_plain_closure_with_move_exprs`, which can wrap the lowered
    // closure in `let` initializers for `move(...)`. Coroutine closures use the
    // same wrapper after building their coroutine-specific body shape.
    pub(super) 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_marker {
            Some(coroutine_marker) => self.lower_expr_coroutine_closure_with_move_exprs(
                expr_hir_id,
                attrs,
                &closure.binder,
                closure.capture_clause,
                e.id,
                coroutine_marker,
                closure.constness,
                &closure.fn_decl,
                &closure.body,
                closure.fn_decl_span,
                closure.fn_arg_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_coroutine_closure_with_move_exprs(
        &mut self,
        expr_hir_id: HirId,
        attrs: &[rustc_attr_ir::Attribute],
        binder: &ClosureBinder,
        capture_clause: CaptureBy,
        closure_id: NodeId,
        coroutine_marker: CoroutineMarker,
        constness: Const,
        decl: &FnDecl,
        body: &Expr,
        fn_decl_span: Span,
        fn_arg_span: Span,
        whole_span: Span,
    ) -> hir::Expr<'hir> {
        let (kind, move_expr_state) =
            self.with_move_expr_bindings(Some(MoveExprState::default()), |this| {
                this.lower_expr_coroutine_closure(
                    binder,
                    capture_clause,
                    closure_id,
                    expr_hir_id,
                    coroutine_marker,
                    constness,
                    decl,
                    body,
                    fn_decl_span,
                    fn_arg_span,
                    attrs,
                )
            });
        let Some(move_expr_state) = move_expr_state else {
            span_bug!(fn_decl_span, "coroutine closure lowering did not return `move(...)` state");
        };
        let closure_expr =
            hir::Expr { hir_id: expr_hir_id, kind, span: self.lower_span(whole_span) };

        self.lower_expr_with_move_exprs(closure_expr, move_expr_state, body, whole_span)
    }

    /// Lowers a plain closure expression and wraps it in an outer block if the
    /// closure body used `move(...)`.
    ///
    /// The lowering is split this way because `move(...)` initializers must be
    /// evaluated before the closure is created, but the closure body must still
    /// lower each `move(...)` occurrence as a use of the synthetic local that
    /// will be introduced by that outer block. For example,
    /// `|| move(foo.clone()).len()` becomes roughly:
    ///
    /// ```ignore (illustrative)
    /// {
    ///     let __move_expr_0 = foo.clone();
    ///     || __move_expr_0.len()
    /// }
    /// ```
    ///
    /// If the initializer contains another `move(...)`, it is lowered after
    /// this closure's state is popped and therefore belongs to the immediately
    /// enclosing closure-like body.
    fn lower_expr_plain_closure_with_move_exprs(
        &mut self,
        expr_hir_id: HirId,
        attrs: &[rustc_attr_ir::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 (closure_kind, move_expr_state) = self.lower_expr_closure(
            attrs,
            binder,
            capture_clause,
            closure_id,
            constness,
            movability,
            decl,
            body,
            fn_decl_span,
            fn_arg_span,
        );

        let closure_expr = hir::Expr {
            hir_id: expr_hir_id,
            kind: closure_kind,
            span: self.lower_span(whole_span),
        };

        self.lower_expr_with_move_exprs(closure_expr, move_expr_state, body, whole_span)
    }

    // Lowers the actual plain closure node and body. The body is lowered while a
    // `MoveExprState` is active, so `move(...)` occurrences become synthetic
    // local uses and the caller can later add the matching initializers.
    fn lower_expr_closure(
        &mut self,
        attrs: &[rustc_attr_ir::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,
    ) -> (hir::ExprKind<'hir>, MoveExprState<'hir>) {
        let closure_def_id = self.local_def_id(closure_id);
        let (binder_clause, generic_params) = self.lower_closure_binder(binder);

        let ((body_id, closure_kind), move_expr_state) =
            self.with_new_scopes(fn_decl_span, move |this| {
                let mut coroutine_kind = find_attr!(
                    attrs,
                    Coroutine => hir::CoroutineKind::Coroutine(Movability::Movable)
                );

                this.with_move_expr_bindings(Some(MoveExprState::default()), |this| {
                    // FIXME(contracts): Support contracts on closures?
                    let body_id = this.lower_fn_body(decl, None, |this| {
                        this.coroutine_kind = coroutine_kind;
                        let e = this.lower_expr_mut(body);
                        coroutine_kind = this.coroutine_kind;
                        e
                    });
                    let coroutine_option = this.closure_movability_for_fn(
                        decl,
                        fn_decl_span,
                        coroutine_kind,
                        movability,
                    );
                    (body_id, coroutine_option)
                })
            });
        let Some(move_expr_state) = move_expr_state else {
            span_bug!(fn_decl_span, "plain closure lowering did not return `move(...)` state");
        };
        let explicit_captures: &'hir [hir::ExplicitCapture] = self.arena.alloc_from_iter(
            move_expr_state
                .occurrences
                .iter()
                .map(|occurrence| hir::ExplicitCapture { var_hir_id: occurrence.binding }),
        );

        let bound_generic_params = self.lower_lifetime_binder(closure_id, generic_params);
        // Lower outside new scope to preserve `is_in_loop_condition`.
        let fn_decl = self.lower_fn_decl(decl, closure_id, FnDeclKind::Closure, None);

        let c = self.arena.alloc(hir::Closure {
            def_id: closure_def_id,
            binder: binder_clause,
            capture_clause: self.lower_capture_clause(capture_clause),
            bound_generic_params,
            fn_decl,
            body: body_id,
            fn_decl_span: self.lower_span(fn_decl_span),
            fn_arg_span: Some(self.lower_span(fn_arg_span)),
            kind: closure_kind,
            constness: self.lower_constness(attrs, constness),
            explicit_captures,
        });

        (hir::ExprKind::Closure(c), move_expr_state)
    }

    fn closure_movability_for_fn(
        &mut self,
        decl: &FnDecl,
        fn_decl_span: Span,
        coroutine_kind: Option<hir::CoroutineKind>,
        movability: Movability,
    ) -> hir::ClosureKind {
        match coroutine_kind {
            Some(hir::CoroutineKind::Coroutine(_)) => {
                if decl.inputs.len() > 1 {
                    self.dcx().emit_err(CoroutineTooManyParameters { fn_decl_span });
                }
                hir::ClosureKind::Coroutine(hir::CoroutineKind::Coroutine(movability))
            }
            Some(
                hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Gen, _)
                | hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Async, _)
                | hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::AsyncGen, _),
            ) => {
                panic!("non-`async`/`gen` closure body turned `async`/`gen` during lowering");
            }
            None => {
                if movability == Movability::Static {
                    self.dcx().emit_err(ClosureCannotBeStatic { fn_decl_span });
                }
                hir::ClosureKind::Closure
            }
        }
    }

    fn lower_closure_binder<'c>(
        &mut self,
        binder: &'c ClosureBinder,
    ) -> (hir::ClosureBinder, &'c [GenericParam]) {
        let (binder, params) = match binder {
            ClosureBinder::NotPresent => (hir::ClosureBinder::Default, &[][..]),
            ClosureBinder::For { span, generic_params } => {
                let span = self.lower_span(*span);
                (hir::ClosureBinder::For { span }, &**generic_params)
            }
        };

        (binder, params)
    }

    // Coroutine closures are lowered separately because they build a different
    // body shape. The source body is lowered with the caller's `MoveExprState`
    // active, so `move(...)` occurrences are collected and hoisted into a block
    // around the outer closure expression.
    fn lower_expr_coroutine_closure(
        &mut self,
        binder: &ClosureBinder,
        capture_clause: CaptureBy,
        closure_id: NodeId,
        closure_hir_id: HirId,
        coroutine_marker: CoroutineMarker,
        constness: Const,
        decl: &FnDecl,
        body: &Expr,
        fn_decl_span: Span,
        fn_arg_span: Span,
        attrs: &[rustc_attr_ir::Attribute],
    ) -> hir::ExprKind<'hir> {
        let closure_def_id = self.local_def_id(closure_id);
        let (binder_clause, generic_params) = self.lower_closure_binder(binder);

        let coroutine_desugaring = match coroutine_marker.kind {
            CoroutineKind::Async => hir::CoroutineDesugaring::Async,
            CoroutineKind::Gen => hir::CoroutineDesugaring::Gen,
            CoroutineKind::AsyncGen => {
                span_bug!(
                    coroutine_marker.span,
                    "only async closures and `iter!` closures are supported currently"
                )
            }
        };

        let body = self.with_new_scopes(fn_decl_span, |this| {
            let inner_decl =
                FnDecl { inputs: decl.inputs.clone(), output: FnRetTy::Default(fn_decl_span) };

            // Transform `async |x: u8| -> X { ... }` into
            // `|x: u8| || -> X { ... }`.
            let body_id = this.lower_body(|this| {
                let (parameters, expr) = this.lower_coroutine_body_with_moved_arguments(
                    &inner_decl,
                    |this| this.with_new_scopes(fn_decl_span, |this| this.lower_expr_mut(body)),
                    fn_decl_span,
                    body.span,
                    coroutine_marker,
                    hir::CoroutineSource::Closure,
                );

                this.maybe_forward_track_caller(closure_hir_id, expr.hir_id);

                (parameters, expr)
            });
            body_id
        });

        let bound_generic_params = self.lower_lifetime_binder(closure_id, generic_params);
        // We need to lower the declaration outside the new scope, because we
        // have to conserve the state of being inside a loop condition for the
        // closure argument types.
        let fn_decl = self.lower_fn_decl(&decl, closure_id, FnDeclKind::Closure, None);

        if let Const::Yes(span) = constness {
            self.dcx().span_err(span, "const coroutines are not supported");
        }

        let explicit_captures: &'hir [hir::ExplicitCapture] = self.arena.alloc_from_iter(
            self.move_expr_bindings
                .last()
                .and_then(Option::as_ref)
                .into_iter()
                .flat_map(|state| &state.occurrences)
                .map(|occurrence| hir::ExplicitCapture { var_hir_id: occurrence.binding }),
        );

        let c = self.arena.alloc(hir::Closure {
            def_id: closure_def_id,
            binder: binder_clause,
            capture_clause: self.lower_capture_clause(capture_clause),
            bound_generic_params,
            fn_decl,
            body,
            fn_decl_span: self.lower_span(fn_decl_span),
            fn_arg_span: Some(self.lower_span(fn_arg_span)),
            // Lower this as a `CoroutineClosure`. That will ensure that HIR typeck
            // 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(attrs, constness),
            explicit_captures,
        });
        hir::ExprKind::Closure(c)
    }
}