From 91e74edca01dbbbb444117f9e85ed76164390314 Mon Sep 17 00:00:00 2001 From: Michael Goulet Date: Sat, 14 Dec 2024 03:45:29 +0000 Subject: [PATCH] Encode coroutine-closures in SMIR --- .../rustc_smir/src/rustc_internal/internal.rs | 4 + compiler/rustc_smir/src/rustc_internal/mod.rs | 4 + .../rustc_smir/src/rustc_smir/convert/mir.rs | 7 +- compiler/stable_mir/src/mir/body.rs | 8 +- compiler/stable_mir/src/mir/pretty.rs | 4 + compiler/stable_mir/src/ty.rs | 11 +++ compiler/stable_mir/src/visitor.rs | 1 + tests/ui/stable-mir-print/async-closure.rs | 12 +++ .../ui/stable-mir-print/async-closure.stdout | 90 +++++++++++++++++++ 9 files changed, 137 insertions(+), 4 deletions(-) create mode 100644 tests/ui/stable-mir-print/async-closure.rs create mode 100644 tests/ui/stable-mir-print/async-closure.stdout diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs index dec2a77619bad..c465367b6b9e0 100644 --- 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 index 614c9169d660b..64d241067a845 100644 --- 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 index fcdf8703b14da..a5a17b4b5730e 100644 --- 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 index f96487cc53c9f..dfd090b39563f 100644 --- 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 index 01a50d46b2d9c..93ed32e258a8b 100644 --- a/compiler/stable_mir/src/mir/pretty.rs +++ b/compiler/stable_mir/src/mir/pretty.rs @@ -410,6 +410,10 @@ fn pretty_aggregate( 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 index 9ce72f155f95c..d7eb435e13f40 100644 --- 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`, 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>, Region, DynKind), Never, Tuple(Vec), @@ -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 index 48260285408c7..3533ed2e8511d 100644 --- 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 new file mode 100644 index 0000000000000..7da532a359f94 --- /dev/null +++ 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 new file mode 100644 index 0000000000000..21df1fd395403 --- /dev/null +++ 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; + } +}