Compute closure captures

Author: HKalbasi

crates/hir-def/src/body/lower.rs

@@ -32,8 +32,8 @@ use crate::{
     builtin_type::{BuiltinFloat, BuiltinInt, BuiltinUint},
     db::DefDatabase,
     expr::{
-        dummy_expr_id, Array, Binding, BindingAnnotation, BindingId, ClosureKind, Expr, ExprId,
-        FloatTypeWrapper, Label, LabelId, Literal, MatchArm, Movability, Pat, PatId,
+        dummy_expr_id, Array, Binding, BindingAnnotation, BindingId, CaptureBy, ClosureKind, Expr,
+        ExprId, FloatTypeWrapper, Label, LabelId, Literal, MatchArm, Movability, Pat, PatId,
         RecordFieldPat, RecordLitField, Statement,
     },
     item_scope::BuiltinShadowMode,
@@ -105,6 +105,7 @@ pub(super) fn lower(
         current_try_block: None,
         is_lowering_assignee_expr: false,
         is_lowering_generator: false,
+        current_binding_owner: None,
     }
     .collect(params, body, is_async_fn)
 }
@@ -119,6 +120,7 @@ struct ExprCollector<'a> {
     current_try_block: Option<LabelId>,
     is_lowering_assignee_expr: bool,
     is_lowering_generator: bool,
+    current_binding_owner: Option<ExprId>,
 }
 
 #[derive(Debug, Default)]
@@ -218,7 +220,12 @@ impl ExprCollector<'_> {
     }
 
     fn alloc_binding(&mut self, name: Name, mode: BindingAnnotation) -> BindingId {
-        self.body.bindings.alloc(Binding { name, mode, definitions: SmallVec::new() })
+        self.body.bindings.alloc(Binding {
+            name,
+            mode,
+            definitions: SmallVec::new(),
+            owner: self.current_binding_owner,
+        })
     }
     fn alloc_pat(&mut self, pat: Pat, ptr: PatPtr) -> PatId {
         let src = self.expander.to_source(ptr);
@@ -497,6 +504,8 @@ impl ExprCollector<'_> {
                 }
             }
             ast::Expr::ClosureExpr(e) => {
+                let (result_expr_id, prev_binding_owner) =
+                    self.initialize_binding_owner(syntax_ptr);
                 let mut args = Vec::new();
                 let mut arg_types = Vec::new();
                 if let Some(pl) = e.param_list() {
@@ -530,18 +539,19 @@ impl ExprCollector<'_> {
                 } else {
                     ClosureKind::Closure
                 };
+                let capture_by =
+                    if e.move_token().is_some() { CaptureBy::Value } else { CaptureBy::Ref };
                 self.is_lowering_generator = prev_is_lowering_generator;
-
-                self.alloc_expr(
-                    Expr::Closure {
-                        args: args.into(),
-                        arg_types: arg_types.into(),
-                        ret_type,
-                        body,
-                        closure_kind,
-                    },
-                    syntax_ptr,
-                )
+                self.current_binding_owner = prev_binding_owner;
+                self.body.exprs[result_expr_id] = Expr::Closure {
+                    args: args.into(),
+                    arg_types: arg_types.into(),
+                    ret_type,
+                    body,
+                    closure_kind,
+                    capture_by,
+                };
+                result_expr_id
             }
             ast::Expr::BinExpr(e) => {
                 let op = e.op_kind();
@@ -581,7 +591,17 @@ impl ExprCollector<'_> {
                     }
                     ArrayExprKind::Repeat { initializer, repeat } => {
                         let initializer = self.collect_expr_opt(initializer);
-                        let repeat = self.collect_expr_opt(repeat);
+                        let repeat = if let Some(repeat) = repeat {
+                            let (id, prev_owner) =
+                                self.initialize_binding_owner(AstPtr::new(&repeat));
+                            let tmp = self.collect_expr(repeat);
+                            self.body.exprs[id] =
+                                mem::replace(&mut self.body.exprs[tmp], Expr::Missing);
+                            self.current_binding_owner = prev_owner;
+                            id
+                        } else {
+                            self.missing_expr()
+                        };
                         self.alloc_expr(
                             Expr::Array(Array::Repeat { initializer, repeat }),
                             syntax_ptr,
@@ -627,6 +647,16 @@ impl ExprCollector<'_> {
         })
     }
 
+    fn initialize_binding_owner(
+        &mut self,
+        syntax_ptr: AstPtr<ast::Expr>,
+    ) -> (ExprId, Option<ExprId>) {
+        let result_expr_id = self.alloc_expr(Expr::Missing, syntax_ptr);
+        let prev_binding_owner = self.current_binding_owner.take();
+        self.current_binding_owner = Some(result_expr_id);
+        (result_expr_id, prev_binding_owner)
+    }
+
     /// Desugar `try { <stmts>; <expr> }` into `'<new_label>: { <stmts>; ::std::ops::Try::from_output(<expr>) }`,
     /// `try { <stmts>; }` into `'<new_label>: { <stmts>; ::std::ops::Try::from_output(()) }`
     /// and save the `<new_label>` to use it as a break target for desugaring of the `?` operator.

crates/hir-def/src/body/pretty.rs

@@ -5,7 +5,9 @@ use std::fmt::{self, Write};
 use syntax::ast::HasName;
 
 use crate::{
-    expr::{Array, BindingAnnotation, BindingId, ClosureKind, Literal, Movability, Statement},
+    expr::{
+        Array, BindingAnnotation, BindingId, CaptureBy, ClosureKind, Literal, Movability, Statement,
+    },
     pretty::{print_generic_args, print_path, print_type_ref},
     type_ref::TypeRef,
 };
@@ -355,7 +357,7 @@ impl<'a> Printer<'a> {
                 self.print_expr(*index);
                 w!(self, "]");
             }
-            Expr::Closure { args, arg_types, ret_type, body, closure_kind } => {
+            Expr::Closure { args, arg_types, ret_type, body, closure_kind, capture_by } => {
                 match closure_kind {
                     ClosureKind::Generator(Movability::Static) => {
                         w!(self, "static ");
@@ -365,6 +367,12 @@ impl<'a> Printer<'a> {
                     }
                     _ => (),
                 }
+                match capture_by {
+                    CaptureBy::Value => {
+                        w!(self, "move ");
+                    }
+                    CaptureBy::Ref => (),
+                }
                 w!(self, "|");
                 for (i, (pat, ty)) in args.iter().zip(arg_types.iter()).enumerate() {
                     if i != 0 {

crates/hir-def/src/expr.rs

@@ -233,6 +233,7 @@ pub enum Expr {
         ret_type: Option<Interned<TypeRef>>,
         body: ExprId,
         closure_kind: ClosureKind,
+        capture_by: CaptureBy,
     },
     Tuple {
         exprs: Box<[ExprId]>,
@@ -250,6 +251,14 @@ pub enum ClosureKind {
     Async,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum CaptureBy {
+    /// `move |x| y + x`.
+    Value,
+    /// `move` keyword was not specified.
+    Ref,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum Movability {
     Static,
@@ -442,6 +451,22 @@ pub struct Binding {
     pub name: Name,
     pub mode: BindingAnnotation,
     pub definitions: SmallVec<[PatId; 1]>,
+    /// Id of the closure/generator that owns this binding. If it is owned by the
+    /// top level expression, this field would be `None`.
+    pub owner: Option<ExprId>,
+}
+
+impl Binding {
+    pub fn is_upvar(&self, relative_to: ExprId) -> bool {
+        match self.owner {
+            Some(x) => {
+                // We assign expression ids in a way that outer closures will recieve
+                // a lower id
+                x.into_raw() < relative_to.into_raw()
+            }
+            None => true,
+        }
+    }
 }
 
 #[derive(Debug, Clone, Eq, PartialEq)]

crates/hir-def/src/macro_expansion_tests/builtin_derive_macro.rs

@@ -16,7 +16,7 @@ struct Foo;
 #[derive(Copy)]
 struct Foo;
 
-impl < > core::marker::Copy for Foo< > {}"#]],
+impl < > core::marker::Copy for Foo< > where {}"#]],
     );
 }
 
@@ -41,7 +41,7 @@ macro Copy {}
 #[derive(Copy)]
 struct Foo;
 
-impl < > crate ::marker::Copy for Foo< > {}"#]],
+impl < > crate ::marker::Copy for Foo< > where {}"#]],
     );
 }
 
@@ -57,7 +57,7 @@ struct Foo<A, B>;
 #[derive(Copy)]
 struct Foo<A, B>;
 
-impl <T0: core::marker::Copy, T1: core::marker::Copy, > core::marker::Copy for Foo<T0, T1, > {}"#]],
+impl <A: core::marker::Copy, B: core::marker::Copy, > core::marker::Copy for Foo<A, B, > where {}"#]],
     );
 }
 
@@ -74,7 +74,7 @@ struct Foo<A, B, 'a, 'b>;
 #[derive(Copy)]
 struct Foo<A, B, 'a, 'b>;
 
-impl <T0: core::marker::Copy, T1: core::marker::Copy, > core::marker::Copy for Foo<T0, T1, > {}"#]],
+impl <A: core::marker::Copy, B: core::marker::Copy, > core::marker::Copy for Foo<A, B, > where {}"#]],
     );
 }
 
@@ -90,7 +90,7 @@ struct Foo<A, B>;
 #[derive(Clone)]
 struct Foo<A, B>;
 
-impl <T0: core::clone::Clone, T1: core::clone::Clone, > core::clone::Clone for Foo<T0, T1, > {}"#]],
+impl <A: core::clone::Clone, B: core::clone::Clone, > core::clone::Clone for Foo<A, B, > where {}"#]],
     );
 }
 
@@ -106,6 +106,6 @@ struct Foo<const X: usize, T>(u32);
 #[derive(Clone)]
 struct Foo<const X: usize, T>(u32);
 
-impl <const T0: usize, T1: core::clone::Clone, > core::clone::Clone for Foo<T0, T1, > {}"#]],
+impl <const X: usize, T: core::clone::Clone, > core::clone::Clone for Foo<X, T, > where u32: core::clone::Clone, {}"#]],
     );
 }

crates/hir-expand/src/builtin_derive_macro.rs

@@ -5,7 +5,7 @@ use tracing::debug;
 
 use crate::tt::{self, TokenId};
 use syntax::{
-    ast::{self, AstNode, HasGenericParams, HasModuleItem, HasName},
+    ast::{self, AstNode, HasGenericParams, HasModuleItem, HasName, HasTypeBounds},
     match_ast,
 };
 
@@ -60,8 +60,11 @@ pub fn find_builtin_derive(ident: &name::Name) -> Option<BuiltinDeriveExpander>
 
 struct BasicAdtInfo {
     name: tt::Ident,
-    /// `Some(ty)` if it's a const param of type `ty`, `None` if it's a type param.
-    param_types: Vec<Option<tt::Subtree>>,
+    /// first field is the name, and
+    /// second field is `Some(ty)` if it's a const param of type `ty`, `None` if it's a type param.
+    /// third fields is where bounds, if any
+    param_types: Vec<(tt::Subtree, Option<tt::Subtree>, Option<tt::Subtree>)>,
+    field_types: Vec<tt::Subtree>,
 }
 
 fn parse_adt(tt: &tt::Subtree) -> Result<BasicAdtInfo, ExpandError> {
@@ -75,17 +78,34 @@ fn parse_adt(tt: &tt::Subtree) -> Result<BasicAdtInfo, ExpandError> {
         ExpandError::Other("no item found".into())
     })?;
     let node = item.syntax();
-    let (name, params) = match_ast! {
+    let (name, params, fields) = match_ast! {
         match node {
-            ast::Struct(it) => (it.name(), it.generic_param_list()),
-            ast::Enum(it) => (it.name(), it.generic_param_list()),
-            ast::Union(it) => (it.name(), it.generic_param_list()),
+            ast::Struct(it) => {
+                (it.name(), it.generic_param_list(), it.field_list())
+            },
+            ast::Enum(it) => (it.name(), it.generic_param_list(), None),
+            ast::Union(it) => (it.name(), it.generic_param_list(), it.record_field_list().map(|x| ast::FieldList::RecordFieldList(x))),
             _ => {
                 debug!("unexpected node is {:?}", node);
                 return Err(ExpandError::Other("expected struct, enum or union".into()))
             },
         }
     };
+    let field_types = match fields {
+        Some(fields) => match fields {
+            ast::FieldList::RecordFieldList(x) => x
+                .fields()
+                .filter_map(|x| x.ty())
+                .map(|x| mbe::syntax_node_to_token_tree(x.syntax()).0)
+                .collect(),
+            ast::FieldList::TupleFieldList(x) => x
+                .fields()
+                .filter_map(|x| x.ty())
+                .map(|x| mbe::syntax_node_to_token_tree(x.syntax()).0)
+                .collect(),
+        },
+        None => vec![],
+    };
     let name = name.ok_or_else(|| {
         debug!("parsed item has no name");
         ExpandError::Other("missing name".into())
@@ -97,35 +117,46 @@ fn parse_adt(tt: &tt::Subtree) -> Result<BasicAdtInfo, ExpandError> {
         .into_iter()
         .flat_map(|param_list| param_list.type_or_const_params())
         .map(|param| {
-            if let ast::TypeOrConstParam::Const(param) = param {
+            let name = param
+                .name()
+                .map(|x| mbe::syntax_node_to_token_tree(x.syntax()).0)
+                .unwrap_or_else(tt::Subtree::empty);
+            let bounds = match &param {
+                ast::TypeOrConstParam::Type(x) => {
+                    x.type_bound_list().map(|x| mbe::syntax_node_to_token_tree(x.syntax()).0)
+                }
+                ast::TypeOrConstParam::Const(_) => None,
+            };
+            let ty = if let ast::TypeOrConstParam::Const(param) = param {
                 let ty = param
                     .ty()
                     .map(|ty| mbe::syntax_node_to_token_tree(ty.syntax()).0)
                     .unwrap_or_else(tt::Subtree::empty);
                 Some(ty)
             } else {
                 None
-            }
+            };
+            (name, ty, bounds)
         })
         .collect();
-    Ok(BasicAdtInfo { name: name_token, param_types })
+    Ok(BasicAdtInfo { name: name_token, param_types, field_types })
 }
 
 fn expand_simple_derive(tt: &tt::Subtree, trait_path: tt::Subtree) -> ExpandResult<tt::Subtree> {
     let info = match parse_adt(tt) {
         Ok(info) => info,
         Err(e) => return ExpandResult::with_err(tt::Subtree::empty(), e),
     };
+    let mut where_block = vec![];
     let (params, args): (Vec<_>, Vec<_>) = info
         .param_types
         .into_iter()
-        .enumerate()
-        .map(|(idx, param_ty)| {
-            let ident = tt::Leaf::Ident(tt::Ident {
-                span: tt::TokenId::unspecified(),
-                text: format!("T{idx}").into(),
-            });
+        .map(|(ident, param_ty, bound)| {
             let ident_ = ident.clone();
+            if let Some(b) = bound {
+                let ident = ident.clone();
+                where_block.push(quote! { #ident : #b , });
+            }
             if let Some(ty) = param_ty {
                 (quote! { const #ident : #ty , }, quote! { #ident_ , })
             } else {
@@ -134,9 +165,16 @@ fn expand_simple_derive(tt: &tt::Subtree, trait_path: tt::Subtree) -> ExpandResu
             }
         })
         .unzip();
+
+    where_block.extend(info.field_types.iter().map(|x| {
+        let x = x.clone();
+        let bound = trait_path.clone();
+        quote! { #x : #bound , }
+    }));
+
     let name = info.name;
     let expanded = quote! {
-        impl < ##params > #trait_path for #name < ##args > {}
+        impl < ##params > #trait_path for #name < ##args > where ##where_block {}
     };
     ExpandResult::ok(expanded)
 }