Rollup merge of #98705 - WaffleLapkin:closure_binder, r=cjgillot

Implement `for<>` lifetime binder for closures

This PR implements RFC 3216 ([TI](#97362)) and allows code like the following:

```rust
let _f = for<'a, 'b> |a: &'a A, b: &'b B| -> &'b C { b.c(a) };
//       ^^^^^^^^^^^--- new!
```

cc ``@Aaron1011`` ``@cjgillot``

Author: Dylan-DPC

compiler/rustc_ast/src/ast.rs

@@ -1390,7 +1390,7 @@ pub enum ExprKind {
     /// A closure (e.g., `move |a, b, c| a + b + c`).
     ///
     /// The final span is the span of the argument block `|...|`.
-    Closure(CaptureBy, Async, Movability, P<FnDecl>, P<Expr>, Span),
+    Closure(ClosureBinder, CaptureBy, Async, Movability, P<FnDecl>, P<Expr>, Span),
     /// A block (`'label: { ... }`).
     Block(P<Block>, Option<Label>),
     /// An async block (`async move { ... }`).
@@ -1518,6 +1518,31 @@ pub enum Movability {
     Movable,
 }
 
+/// Closure lifetime binder, `for<'a, 'b>` in `for<'a, 'b> |_: &'a (), _: &'b ()|`.
+#[derive(Clone, Encodable, Decodable, Debug)]
+pub enum ClosureBinder {
+    /// The binder is not present, all closure lifetimes are inferred.
+    NotPresent,
+    /// The binder is present.
+    For {
+        /// Span of the whole `for<>` clause
+        ///
+        /// ```text
+        /// for<'a, 'b> |_: &'a (), _: &'b ()| { ... }
+        /// ^^^^^^^^^^^ -- this
+        /// ```
+        span: Span,
+
+        /// Lifetimes in the `for<>` closure
+        ///
+        /// ```text
+        /// for<'a, 'b> |_: &'a (), _: &'b ()| { ... }
+        ///     ^^^^^^ -- this
+        /// ```
+        generic_params: P<[GenericParam]>,
+    },
+}
+
 /// Represents a macro invocation. The `path` indicates which macro
 /// is being invoked, and the `args` are arguments passed to it.
 #[derive(Clone, Encodable, Decodable, Debug)]

compiler/rustc_ast/src/mut_visit.rs

@@ -125,6 +125,10 @@ pub trait MutVisitor: Sized {
         noop_visit_asyncness(a, self);
     }
 
+    fn visit_closure_binder(&mut self, b: &mut ClosureBinder) {
+        noop_visit_closure_binder(b, self);
+    }
+
     fn visit_block(&mut self, b: &mut P<Block>) {
         noop_visit_block(b, self);
     }
@@ -825,6 +829,17 @@ pub fn visit_constness<T: MutVisitor>(constness: &mut Const, vis: &mut T) {
     }
 }
 
+pub fn noop_visit_closure_binder<T: MutVisitor>(binder: &mut ClosureBinder, vis: &mut T) {
+    match binder {
+        ClosureBinder::NotPresent => {}
+        ClosureBinder::For { span: _, generic_params } => {
+            let mut vec = std::mem::take(generic_params).into_vec();
+            vec.flat_map_in_place(|param| vis.flat_map_generic_param(param));
+            *generic_params = P::from_vec(vec);
+        }
+    }
+}
+
 pub fn noop_visit_asyncness<T: MutVisitor>(asyncness: &mut Async, vis: &mut T) {
     match asyncness {
         Async::Yes { span: _, closure_id, return_impl_trait_id } => {
@@ -1336,7 +1351,8 @@ pub fn noop_visit_expr<T: MutVisitor>(
             vis.visit_expr(expr);
             arms.flat_map_in_place(|arm| vis.flat_map_arm(arm));
         }
-        ExprKind::Closure(_capture_by, asyncness, _movability, decl, body, span) => {
+        ExprKind::Closure(binder, _capture_by, asyncness, _movability, decl, body, span) => {
+            vis.visit_closure_binder(binder);
             vis.visit_asyncness(asyncness);
             vis.visit_fn_decl(decl);
             vis.visit_expr(body);

compiler/rustc_ast/src/visit.rs

@@ -56,14 +56,14 @@ pub enum FnKind<'a> {
     Fn(FnCtxt, Ident, &'a FnSig, &'a Visibility, &'a Generics, Option<&'a Block>),
 
     /// E.g., `|x, y| body`.
-    Closure(&'a FnDecl, &'a Expr),
+    Closure(&'a ClosureBinder, &'a FnDecl, &'a Expr),
 }
 
 impl<'a> FnKind<'a> {
     pub fn header(&self) -> Option<&'a FnHeader> {
         match *self {
             FnKind::Fn(_, _, sig, _, _, _) => Some(&sig.header),
-            FnKind::Closure(_, _) => None,
+            FnKind::Closure(_, _, _) => None,
         }
     }
 
@@ -77,7 +77,7 @@ impl<'a> FnKind<'a> {
     pub fn decl(&self) -> &'a FnDecl {
         match self {
             FnKind::Fn(_, _, sig, _, _, _) => &sig.decl,
-            FnKind::Closure(decl, _) => decl,
+            FnKind::Closure(_, decl, _) => decl,
         }
     }
 
@@ -155,6 +155,9 @@ pub trait Visitor<'ast>: Sized {
     fn visit_generics(&mut self, g: &'ast Generics) {
         walk_generics(self, g)
     }
+    fn visit_closure_binder(&mut self, b: &'ast ClosureBinder) {
+        walk_closure_binder(self, b)
+    }
     fn visit_where_predicate(&mut self, p: &'ast WherePredicate) {
         walk_where_predicate(self, p)
     }
@@ -636,6 +639,15 @@ pub fn walk_generics<'a, V: Visitor<'a>>(visitor: &mut V, generics: &'a Generics
     walk_list!(visitor, visit_where_predicate, &generics.where_clause.predicates);
 }
 
+pub fn walk_closure_binder<'a, V: Visitor<'a>>(visitor: &mut V, binder: &'a ClosureBinder) {
+    match binder {
+        ClosureBinder::NotPresent => {}
+        ClosureBinder::For { generic_params, span: _ } => {
+            walk_list!(visitor, visit_generic_param, generic_params)
+        }
+    }
+}
+
 pub fn walk_where_predicate<'a, V: Visitor<'a>>(visitor: &mut V, predicate: &'a WherePredicate) {
     match *predicate {
         WherePredicate::BoundPredicate(WhereBoundPredicate {
@@ -682,7 +694,8 @@ pub fn walk_fn<'a, V: Visitor<'a>>(visitor: &mut V, kind: FnKind<'a>, _span: Spa
             walk_fn_decl(visitor, &sig.decl);
             walk_list!(visitor, visit_block, body);
         }
-        FnKind::Closure(decl, body) => {
+        FnKind::Closure(binder, decl, body) => {
+            visitor.visit_closure_binder(binder);
             walk_fn_decl(visitor, decl);
             visitor.visit_expr(body);
         }
@@ -856,8 +869,8 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
             visitor.visit_expr(subexpression);
             walk_list!(visitor, visit_arm, arms);
         }
-        ExprKind::Closure(_, _, _, ref decl, ref body, _decl_span) => {
-            visitor.visit_fn(FnKind::Closure(decl, body), expression.span, expression.id)
+        ExprKind::Closure(ref binder, _, _, _, ref decl, ref body, _decl_span) => {
+            visitor.visit_fn(FnKind::Closure(binder, decl, body), expression.span, expression.id)
         }
         ExprKind::Block(ref block, ref opt_label) => {
             walk_list!(visitor, visit_label, opt_label);

compiler/rustc_ast_lowering/src/expr.rs

@@ -155,6 +155,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     self.lower_expr_await(span, expr)
                 }
                 ExprKind::Closure(
+                    ref binder,
                     capture_clause,
                     asyncness,
                     movability,
@@ -164,6 +165,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 ) => {
                     if let Async::Yes { closure_id, .. } = asyncness {
                         self.lower_expr_async_closure(
+                            binder,
                             capture_clause,
                             e.id,
                             closure_id,
@@ -173,6 +175,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         )
                     } else {
                         self.lower_expr_closure(
+                            binder,
                             capture_clause,
                             e.id,
                             movability,
@@ -605,13 +608,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
         });
 
         // `static |_task_context| -> <ret_ty> { body }`:
-        let generator_kind = hir::ExprKind::Closure {
-            capture_clause,
-            bound_generic_params: &[],
-            fn_decl,
-            body,
-            fn_decl_span: self.lower_span(span),
-            movability: Some(hir::Movability::Static),
+        let generator_kind = {
+            let c = self.arena.alloc(hir::Closure {
+                binder: hir::ClosureBinder::Default,
+                capture_clause,
+                bound_generic_params: &[],
+                fn_decl,
+                body,
+                fn_decl_span: self.lower_span(span),
+                movability: Some(hir::Movability::Static),
+            });
+
+            hir::ExprKind::Closure(c)
         };
         let generator = hir::Expr {
             hir_id: self.lower_node_id(closure_node_id),
@@ -831,14 +839,17 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
     fn lower_expr_closure(
         &mut self,
+        binder: &ClosureBinder,
         capture_clause: CaptureBy,
         closure_id: NodeId,
         movability: Movability,
         decl: &FnDecl,
         body: &Expr,
         fn_decl_span: Span,
     ) -> hir::ExprKind<'hir> {
-        let (body, generator_option) = self.with_new_scopes(move |this| {
+        let (binder_clause, generic_params) = self.lower_closure_binder(binder);
+
+        let (body_id, generator_option) = self.with_new_scopes(move |this| {
             let prev = this.current_item;
             this.current_item = Some(fn_decl_span);
             let mut generator_kind = None;
@@ -853,18 +864,21 @@ impl<'hir> LoweringContext<'_, 'hir> {
             (body_id, generator_option)
         });
 
-        self.with_lifetime_binder(closure_id, &[], |this, bound_generic_params| {
+        self.with_lifetime_binder(closure_id, generic_params, |this, bound_generic_params| {
             // Lower outside new scope to preserve `is_in_loop_condition`.
             let fn_decl = this.lower_fn_decl(decl, None, FnDeclKind::Closure, None);
 
-            hir::ExprKind::Closure {
+            let c = self.arena.alloc(hir::Closure {
+                binder: binder_clause,
                 capture_clause,
                 bound_generic_params,
                 fn_decl,
-                body,
+                body: body_id,
                 fn_decl_span: this.lower_span(fn_decl_span),
                 movability: generator_option,
-            }
+            });
+
+            hir::ExprKind::Closure(c)
         })
     }
 
@@ -906,15 +920,40 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }
     }
 
+    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, ref generic_params } => {
+                let span = self.lower_span(span);
+                (hir::ClosureBinder::For { span }, &**generic_params)
+            }
+        };
+
+        (binder, params)
+    }
+
     fn lower_expr_async_closure(
         &mut self,
+        binder: &ClosureBinder,
         capture_clause: CaptureBy,
         closure_id: NodeId,
         inner_closure_id: NodeId,
         decl: &FnDecl,
         body: &Expr,
         fn_decl_span: Span,
     ) -> hir::ExprKind<'hir> {
+        if let &ClosureBinder::For { span, .. } = binder {
+            self.tcx.sess.span_err(
+                span,
+                "`for<...>` binders on `async` closures are not currently supported",
+            );
+        }
+
+        let (binder_clause, generic_params) = self.lower_closure_binder(binder);
+
         let outer_decl =
             FnDecl { inputs: decl.inputs.clone(), output: FnRetTy::Default(fn_decl_span) };
 
@@ -952,20 +991,22 @@ impl<'hir> LoweringContext<'_, 'hir> {
             body_id
         });
 
-        self.with_lifetime_binder(closure_id, &[], |this, bound_generic_params| {
+        self.with_lifetime_binder(closure_id, generic_params, |this, bound_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 = this.lower_fn_decl(&outer_decl, None, FnDeclKind::Closure, None);
 
-            hir::ExprKind::Closure {
+            let c = self.arena.alloc(hir::Closure {
+                binder: binder_clause,
                 capture_clause,
                 bound_generic_params,
                 fn_decl,
                 body,
                 fn_decl_span: this.lower_span(fn_decl_span),
                 movability: None,
-            }
+            });
+            hir::ExprKind::Closure(c)
         })
     }
 

compiler/rustc_ast_passes/src/ast_validation.rs

@@ -1597,6 +1597,10 @@ impl<'a> Visitor<'a> for AstValidator<'a> {
                 .emit();
         }
 
+        if let FnKind::Closure(ClosureBinder::For { generic_params, .. }, ..) = fk {
+            self.check_late_bound_lifetime_defs(generic_params);
+        }
+
         if let FnKind::Fn(
             _,
             _,

compiler/rustc_ast_passes/src/feature_gate.rs

@@ -744,6 +744,11 @@ pub fn check_crate(krate: &ast::Crate, sess: &Session) {
         "async closures are unstable",
         "to use an async block, remove the `||`: `async {`"
     );
+    gate_all!(
+        closure_lifetime_binder,
+        "`for<...>` binders for closures are experimental",
+        "consider removing `for<...>`"
+    );
     gate_all!(more_qualified_paths, "usage of qualified paths in this context is experimental");
     gate_all!(generators, "yield syntax is experimental");
     gate_all!(raw_ref_op, "raw address of syntax is experimental");

compiler/rustc_ast_pretty/src/pprust/state/expr.rs

@@ -389,13 +389,15 @@ impl<'a> State<'a> {
                 self.bclose(expr.span, empty);
             }
             ast::ExprKind::Closure(
+                ref binder,
                 capture_clause,
                 asyncness,
                 movability,
                 ref decl,
                 ref body,
                 _,
             ) => {
+                self.print_closure_binder(binder);
                 self.print_movability(movability);
                 self.print_asyncness(asyncness);
                 self.print_capture_clause(capture_clause);
@@ -594,6 +596,15 @@ impl<'a> State<'a> {
         self.end(); // Close enclosing cbox.
     }
 
+    fn print_closure_binder(&mut self, binder: &ast::ClosureBinder) {
+        match binder {
+            ast::ClosureBinder::NotPresent => {}
+            ast::ClosureBinder::For { generic_params, .. } => {
+                self.print_formal_generic_params(&generic_params)
+            }
+        }
+    }
+
     fn print_movability(&mut self, movability: ast::Movability) {
         match movability {
             ast::Movability::Static => self.word_space("static"),

compiler/rustc_borrowck/src/diagnostics/mod.rs

@@ -891,7 +891,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         let hir_id = self.infcx.tcx.hir().local_def_id_to_hir_id(local_did);
         let expr = &self.infcx.tcx.hir().expect_expr(hir_id).kind;
         debug!("closure_span: hir_id={:?} expr={:?}", hir_id, expr);
-        if let hir::ExprKind::Closure { body, fn_decl_span, .. } = expr {
+        if let hir::ExprKind::Closure(&hir::Closure { body, fn_decl_span, .. }) = expr {
             for (captured_place, place) in self
                 .infcx
                 .tcx
@@ -904,11 +904,11 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                         if target_place == place.as_ref() =>
                     {
                         debug!("closure_span: found captured local {:?}", place);
-                        let body = self.infcx.tcx.hir().body(*body);
+                        let body = self.infcx.tcx.hir().body(body);
                         let generator_kind = body.generator_kind();
 
                         return Some((
-                            *fn_decl_span,
+                            fn_decl_span,
                             generator_kind,
                             captured_place.get_capture_kind_span(self.infcx.tcx),
                             captured_place.get_path_span(self.infcx.tcx),