From ba64c93a44c5bdd78b01ef868435dc427f0cf4d4 Mon Sep 17 00:00:00 2001
From: Ryo Yoshida <low.ryoshida@gmail.com>
Date: Mon, 5 Sep 2022 22:43:26 +0900
Subject: [PATCH] Lower generator expression to HIR

---
 crates/hir-def/src/body/lower.rs  | 26 ++++++++++++++++++++++++--
 crates/hir-def/src/body/pretty.rs |  7 +++++--
 crates/hir-def/src/expr.rs        | 13 +++++++++++++
 crates/hir-ty/src/infer/expr.rs   |  2 +-
 4 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/crates/hir-def/src/body/lower.rs b/crates/hir-def/src/body/lower.rs
index 3b3297f7811c..c4f91e49a6e1 100644
--- a/crates/hir-def/src/body/lower.rs
+++ b/crates/hir-def/src/body/lower.rs
@@ -29,8 +29,9 @@ use crate::{
     builtin_type::{BuiltinFloat, BuiltinInt, BuiltinUint},
     db::DefDatabase,
     expr::{
-        dummy_expr_id, Array, BindingAnnotation, Expr, ExprId, FloatTypeWrapper, Label, LabelId,
-        Literal, MatchArm, Pat, PatId, RecordFieldPat, RecordLitField, Statement,
+        dummy_expr_id, Array, BindingAnnotation, ClosureKind, Expr, ExprId, FloatTypeWrapper,
+        Label, LabelId, Literal, MatchArm, Movability, Pat, PatId, RecordFieldPat, RecordLitField,
+        Statement,
     },
     intern::Interned,
     item_scope::BuiltinShadowMode,
@@ -97,6 +98,7 @@ pub(super) fn lower(
         name_to_pat_grouping: Default::default(),
         is_lowering_inside_or_pat: false,
         is_lowering_assignee_expr: false,
+        is_lowering_generator: false,
     }
     .collect(params, body)
 }
@@ -111,6 +113,7 @@ struct ExprCollector<'a> {
     name_to_pat_grouping: FxHashMap<Name, Vec<PatId>>,
     is_lowering_inside_or_pat: bool,
     is_lowering_assignee_expr: bool,
+    is_lowering_generator: bool,
 }
 
 impl ExprCollector<'_> {
@@ -358,6 +361,7 @@ impl ExprCollector<'_> {
                 self.alloc_expr(Expr::Return { expr }, syntax_ptr)
             }
             ast::Expr::YieldExpr(e) => {
+                self.is_lowering_generator = true;
                 let expr = e.expr().map(|e| self.collect_expr(e));
                 self.alloc_expr(Expr::Yield { expr }, syntax_ptr)
             }
@@ -459,13 +463,31 @@ impl ExprCollector<'_> {
                     .ret_type()
                     .and_then(|r| r.ty())
                     .map(|it| Interned::new(TypeRef::from_ast(&self.ctx(), it)));
+
+                let prev_is_lowering_generator = self.is_lowering_generator;
+                self.is_lowering_generator = false;
+
                 let body = self.collect_expr_opt(e.body());
+
+                let closure_kind = if self.is_lowering_generator {
+                    let movability = if e.static_token().is_some() {
+                        Movability::Static
+                    } else {
+                        Movability::Movable
+                    };
+                    ClosureKind::Generator(movability)
+                } else {
+                    ClosureKind::Closure
+                };
+                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,
                 )
diff --git a/crates/hir-def/src/body/pretty.rs b/crates/hir-def/src/body/pretty.rs
index f2fed954444e..35686af38854 100644
--- a/crates/hir-def/src/body/pretty.rs
+++ b/crates/hir-def/src/body/pretty.rs
@@ -3,7 +3,7 @@
 use std::fmt::{self, Write};
 
 use crate::{
-    expr::{Array, BindingAnnotation, Literal, Statement},
+    expr::{Array, BindingAnnotation, ClosureKind, Literal, Movability, Statement},
     pretty::{print_generic_args, print_path, print_type_ref},
     type_ref::TypeRef,
 };
@@ -350,7 +350,10 @@ impl<'a> Printer<'a> {
                 self.print_expr(*index);
                 w!(self, "]");
             }
-            Expr::Closure { args, arg_types, ret_type, body } => {
+            Expr::Closure { args, arg_types, ret_type, body, closure_kind } => {
+                if let ClosureKind::Generator(Movability::Static) = closure_kind {
+                    w!(self, "static ");
+                }
                 w!(self, "|");
                 for (i, (pat, ty)) in args.iter().zip(arg_types.iter()).enumerate() {
                     if i != 0 {
diff --git a/crates/hir-def/src/expr.rs b/crates/hir-def/src/expr.rs
index 419d3feec3b6..162646550207 100644
--- a/crates/hir-def/src/expr.rs
+++ b/crates/hir-def/src/expr.rs
@@ -198,6 +198,7 @@ pub enum Expr {
         arg_types: Box<[Option<Interned<TypeRef>>]>,
         ret_type: Option<Interned<TypeRef>>,
         body: ExprId,
+        closure_kind: ClosureKind,
     },
     Tuple {
         exprs: Box<[ExprId]>,
@@ -211,6 +212,18 @@ pub enum Expr {
     Underscore,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum ClosureKind {
+    Closure,
+    Generator(Movability),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum Movability {
+    Static,
+    Movable,
+}
+
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub enum Array {
     ElementList { elements: Box<[ExprId]>, is_assignee_expr: bool },
diff --git a/crates/hir-ty/src/infer/expr.rs b/crates/hir-ty/src/infer/expr.rs
index 2d04a864a2cf..f0382846a6af 100644
--- a/crates/hir-ty/src/infer/expr.rs
+++ b/crates/hir-ty/src/infer/expr.rs
@@ -216,7 +216,7 @@ impl<'a> InferenceContext<'a> {
                 self.diverges = Diverges::Maybe;
                 TyBuilder::unit()
             }
-            Expr::Closure { body, args, ret_type, arg_types } => {
+            Expr::Closure { body, args, ret_type, arg_types, closure_kind: _ } => {
                 assert_eq!(args.len(), arg_types.len());
 
                 let mut sig_tys = Vec::new();