Encode coroutine-closures in SMIR

rust-lang · Dec 14, 2024 · 91e74ed · 91e74ed
1 parent 4a204be
commit 91e74ed
Show file tree

Hide file tree

Showing 9 changed files with 137 additions and 4 deletions.
diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs
@@ -141,6 +141,10 @@ impl RustcInternal for RigidTy {
             RigidTy::Coroutine(def, args, _mov) => {
                 rustc_ty::TyKind::Coroutine(def.0.internal(tables, tcx), args.internal(tables, tcx))
             }
+            RigidTy::CoroutineClosure(def, args) => rustc_ty::TyKind::CoroutineClosure(
+                def.0.internal(tables, tcx),
+                args.internal(tables, tcx),
+            ),
             RigidTy::CoroutineWitness(def, args) => rustc_ty::TyKind::CoroutineWitness(
                 def.0.internal(tables, tcx),
                 args.internal(tables, tcx),

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -107,6 +107,10 @@ impl<'tcx> Tables<'tcx> {
         stable_mir::ty::CoroutineDef(self.create_def_id(did))
     }
 
+    pub fn coroutine_closure_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineClosureDef {
+        stable_mir::ty::CoroutineClosureDef(self.create_def_id(did))
+    }
+
     pub fn alias_def(&mut self, did: DefId) -> stable_mir::ty::AliasDef {
         stable_mir::ty::AliasDef(self.create_def_id(did))
     }

diff --git a/compiler/rustc_smir/src/rustc_smir/convert/mir.rs b/compiler/rustc_smir/src/rustc_smir/convert/mir.rs
@@ -565,8 +565,11 @@ impl<'tcx> Stable<'tcx> for mir::AggregateKind<'tcx> {
                     tables.tcx.coroutine_movability(*def_id).stable(tables),
                 )
             }
-            mir::AggregateKind::CoroutineClosure(..) => {
-                todo!("FIXME(async_closures): Lower these to SMIR")
+            mir::AggregateKind::CoroutineClosure(def_id, generic_args) => {
+                stable_mir::mir::AggregateKind::CoroutineClosure(
+                    tables.coroutine_closure_def(*def_id),
+                    generic_args.stable(tables),
+                )
             }
             mir::AggregateKind::RawPtr(ty, mutability) => {
                 stable_mir::mir::AggregateKind::RawPtr(ty.stable(tables), mutability.stable(tables))

diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
@@ -5,8 +5,8 @@ use serde::Serialize;
 use crate::compiler_interface::with;
 use crate::mir::pretty::function_body;
 use crate::ty::{
-    AdtDef, ClosureDef, CoroutineDef, GenericArgs, MirConst, Movability, Region, RigidTy, Ty,
-    TyConst, TyKind, VariantIdx,
+    AdtDef, ClosureDef, CoroutineClosureDef, CoroutineDef, GenericArgs, MirConst, Movability,
+    Region, RigidTy, Ty, TyConst, TyKind, VariantIdx,
 };
 use crate::{Error, Opaque, Span, Symbol};
 
@@ -617,6 +617,9 @@ impl Rvalue {
                 AggregateKind::Coroutine(def, ref args, mov) => {
                     Ok(Ty::new_coroutine(def, args.clone(), mov))
                 }
+                AggregateKind::CoroutineClosure(def, ref args) => {
+                    Ok(Ty::new_coroutine_closure(def, args.clone()))
+                }
                 AggregateKind::RawPtr(ty, mutability) => Ok(Ty::new_ptr(ty, mutability)),
             },
             Rvalue::ShallowInitBox(_, ty) => Ok(Ty::new_box(*ty)),
@@ -633,6 +636,7 @@ pub enum AggregateKind {
     Closure(ClosureDef, GenericArgs),
     // FIXME(stable_mir): Movability here is redundant
     Coroutine(CoroutineDef, GenericArgs, Movability),
+    CoroutineClosure(CoroutineClosureDef, GenericArgs),
     RawPtr(Ty, Mutability),
 }
 

diff --git a/compiler/stable_mir/src/mir/pretty.rs b/compiler/stable_mir/src/mir/pretty.rs
@@ -410,6 +410,10 @@ fn pretty_aggregate<W: Write>(
             write!(writer, "{{coroutine@{}}}(", def.span().diagnostic())?;
             ")"
         }
+        AggregateKind::CoroutineClosure(def, _) => {
+            write!(writer, "{{coroutine-closure@{}}}(", def.span().diagnostic())?;
+            ")"
+        }
         AggregateKind::RawPtr(ty, mutability) => {
             write!(
                 writer,

diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
@@ -63,6 +63,11 @@ impl Ty {
         Ty::from_rigid_kind(RigidTy::Coroutine(def, args, mov))
     }
 
+    /// Create a new closure type.
+    pub fn new_coroutine_closure(def: CoroutineClosureDef, args: GenericArgs) -> Ty {
+        Ty::from_rigid_kind(RigidTy::CoroutineClosure(def, args))
+    }
+
     /// Create a new box type that represents `Box<T>`, for the given inner type `T`.
     pub fn new_box(inner_ty: Ty) -> Ty {
         with(|cx| cx.new_box_ty(inner_ty))
@@ -550,6 +555,7 @@ pub enum RigidTy {
     Closure(ClosureDef, GenericArgs),
     // FIXME(stable_mir): Movability here is redundant
     Coroutine(CoroutineDef, GenericArgs, Movability),
+    CoroutineClosure(CoroutineClosureDef, GenericArgs),
     Dynamic(Vec<Binder<ExistentialPredicate>>, Region, DynKind),
     Never,
     Tuple(Vec<Ty>),
@@ -740,6 +746,11 @@ crate_def! {
     pub CoroutineDef;
 }
 
+crate_def! {
+    #[derive(Serialize)]
+    pub CoroutineClosureDef;
+}
+
 crate_def! {
     #[derive(Serialize)]
     pub ParamDef;

diff --git a/compiler/stable_mir/src/visitor.rs b/compiler/stable_mir/src/visitor.rs
@@ -168,6 +168,7 @@ impl Visitable for RigidTy {
             | RigidTy::Closure(_, args)
             | RigidTy::Coroutine(_, args, _)
             | RigidTy::CoroutineWitness(_, args)
+            | RigidTy::CoroutineClosure(_, args)
             | RigidTy::FnDef(_, args) => args.visit(visitor),
             RigidTy::FnPtr(sig) => sig.visit(visitor),
             RigidTy::Dynamic(pred, r, _) => {

diff --git a/tests/ui/stable-mir-print/async-closure.rs b/tests/ui/stable-mir-print/async-closure.rs
@@ -0,0 +1,12 @@
+//@ compile-flags: -Z unpretty=stable-mir --crate-type lib -C panic=abort
+//@ check-pass
+//@ only-x86_64
+//@ edition: 2024
+//@ needs-unwind unwind edges are different with panic=abort
+
+pub fn foo() {
+    let y = 0;
+    let x = async || {
+        let y = y;
+    };
+}
diff --git a/tests/ui/stable-mir-print/async-closure.stdout b/tests/ui/stable-mir-print/async-closure.stdout
@@ -0,0 +1,90 @@
+// WARNING: This is highly experimental output it's intended for stable-mir developers only.
+// If you find a bug or want to improve the output open a issue at https://github.com/rust-lang/project-stable-mir.
+fn foo() -> () {
+    let mut _0: ();
+    let  _1: i32;
+    let  _2: {async closure@$DIR/async-closure.rs:9:13: 9:21};
+    let mut _3: &i32;
+    debug y => _1;
+    debug x => _2;
+    bb0: {
+        _1 = 0_i32;
+        _3 = &_1;
+        _2 = {coroutine-closure@$DIR/async-closure.rs:9:13: 9:21}(move _3);
+        return;
+    }
+}
+fn foo::{closure#0}(_1: &{async closure@$DIR/async-closure.rs:9:13: 9:21}) -> {async closure body@$DIR/async-closure.rs:9:22: 11:6} {
+    let mut _0: {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    let mut _2: &i32;
+    debug y => (*((*_1).0: &i32));
+    bb0: {
+        _2 = CopyForDeref(((*_1).0: &i32));
+        _0 = {coroutine@$DIR/async-closure.rs:9:22: 11:6}(_2);
+        return;
+    }
+}
+fn foo::{closure#0}::{closure#0}(_1: Pin<&mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}>, _2: &mut Context<'_>) -> Poll<()> {
+    let mut _0: Poll<()>;
+    let  _3: i32;
+    let mut _4: &i32;
+    let mut _5: u32;
+    let mut _6: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    let mut _7: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    let mut _8: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    debug _task_context => _2;
+    debug y => (*((*(_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6})).0: &i32));
+    debug y => _3;
+    bb0: {
+        _6 = CopyForDeref((_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}));
+        _5 = discriminant((*_6));
+        switchInt(move _5) -> [0: bb1, 1: bb2, otherwise: bb3];
+    }
+    bb1: {
+        _7 = CopyForDeref((_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}));
+        _4 = CopyForDeref(((*_7).0: &i32));
+        _3 = (*_4);
+        _0 = std::task::Poll::Ready(());
+        _8 = CopyForDeref((_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}));
+        discriminant((*_8) = 1;
+        return;
+    }
+    bb2: {
+        assert(false, `async fn` resumed after completion) -> [success: bb2, unwind unreachable];
+    }
+    bb3: {
+        unreachable;
+    }
+}
+fn foo::{closure#0}::{closure#1}(_1: Pin<&mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}>, _2: &mut Context<'_>) -> Poll<()> {
+    let mut _0: Poll<()>;
+    let  _3: i32;
+    let mut _4: &i32;
+    let mut _5: u32;
+    let mut _6: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    let mut _7: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    let mut _8: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6};
+    debug _task_context => _2;
+    debug y => (*((*(_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6})).0: &i32));
+    debug y => _3;
+    bb0: {
+        _6 = CopyForDeref((_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}));
+        _5 = discriminant((*_6));
+        switchInt(move _5) -> [0: bb1, 1: bb2, otherwise: bb3];
+    }
+    bb1: {
+        _7 = CopyForDeref((_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}));
+        _4 = CopyForDeref(((*_7).0: &i32));
+        _3 = (*_4);
+        _0 = std::task::Poll::Ready(());
+        _8 = CopyForDeref((_1.0: &mut {async closure body@$DIR/async-closure.rs:9:22: 11:6}));
+        discriminant((*_8) = 1;
+        return;
+    }
+    bb2: {
+        assert(false, `async fn` resumed after completion) -> [success: bb2, unwind unreachable];
+    }
+    bb3: {
+        unreachable;
+    }
+}