diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs index e331e5934716a..c6ff7c432b9a3 100644 --- a/compiler/rustc_smir/src/stable_mir/ty.rs +++ b/compiler/rustc_smir/src/stable_mir/ty.rs @@ -742,6 +742,14 @@ crate_def! { pub ClosureDef; } +impl ClosureDef { + /// Retrieves the body of the closure definition. Returns None if the body + /// isn't available. + pub fn body(&self) -> Option { + with(|ctx| ctx.has_body(self.0).then(|| ctx.mir_body(self.0))) + } +} + crate_def! { #[derive(Serialize)] pub CoroutineDef; diff --git a/tests/ui-fulldeps/stable-mir/closure-generic-body.rs b/tests/ui-fulldeps/stable-mir/closure-generic-body.rs new file mode 100644 index 0000000000000..2a23345a9d315 --- /dev/null +++ b/tests/ui-fulldeps/stable-mir/closure-generic-body.rs @@ -0,0 +1,90 @@ +//@ run-pass +//! Tests stable mir API for retrieving the body of a closure. + +//@ ignore-stage1 +//@ ignore-cross-compile +//@ ignore-remote +//@ edition: 2021 + +#![feature(rustc_private)] +#![feature(assert_matches)] + +extern crate rustc_middle; +#[macro_use] +extern crate rustc_smir; +extern crate rustc_driver; +extern crate rustc_interface; +extern crate stable_mir; + +use std::io::Write; +use std::ops::ControlFlow; + +use stable_mir::mir::{Body, ConstOperand, Operand, TerminatorKind}; +use stable_mir::ty::{FnDef, RigidTy, TyKind}; + +const CRATE_NAME: &str = "crate_closure_body"; + +fn test_closure_body() -> ControlFlow<()> { + let crate_items = stable_mir::all_local_items(); + for item in crate_items { + let item_ty = item.ty(); + match &item_ty.kind() { + TyKind::RigidTy(RigidTy::Closure(closure_def, _)) => { + let closure_body = closure_def.body().unwrap(); + check_incr_closure_body(closure_body); + } + _ => {} + } + } + + ControlFlow::Continue(()) +} + +fn check_incr_closure_body(body: Body) { + let first_block = &body.blocks[0]; + let TerminatorKind::Call { func: Operand::Constant(ConstOperand { const_, .. }), .. } = + &first_block.terminator.kind + else { + panic!("expected Call Terminator, got: "); + }; + + let TyKind::RigidTy(RigidTy::FnDef(FnDef(def_id), ..), ..) = const_.ty().kind() else { + panic!("expected FnDef"); + }; + + assert_eq!(def_id.name(), "id"); +} + +fn main() { + let path = "closure_body.rs"; + generate_input(&path).unwrap(); + let args = &[ + "rustc".to_string(), + "-Cpanic=abort".to_string(), + "--crate-name".to_string(), + CRATE_NAME.to_string(), + path.to_string(), + ]; + run!(args, test_closure_body).unwrap(); +} + +fn generate_input(path: &str) -> std::io::Result<()> { + let mut file = std::fs::File::create(path)?; + write!( + file, + r#" + fn id(y: T) -> T {{ + y + }} + + fn main() {{ + let cl_id= |x| {{ + id(x) + }}; + + let _= cl_id(5); + }} + "# + )?; + Ok(()) +} diff --git a/tests/ui-fulldeps/stable-mir/closure_body.rs b/tests/ui-fulldeps/stable-mir/closure_body.rs new file mode 100644 index 0000000000000..7ed0dabd2c3e2 --- /dev/null +++ b/tests/ui-fulldeps/stable-mir/closure_body.rs @@ -0,0 +1,90 @@ +//@ run-pass +//! Tests stable mir API for retrieving the body of a closure. + +//@ ignore-stage1 +//@ ignore-cross-compile +//@ ignore-remote +//@ edition: 2021 + +#![feature(rustc_private)] +#![feature(assert_matches)] + +extern crate rustc_middle; +#[macro_use] +extern crate rustc_smir; +extern crate rustc_driver; +extern crate rustc_interface; +extern crate stable_mir; + +use std::io::Write; +use std::ops::ControlFlow; + +use stable_mir::mir::{Body, ConstOperand, Operand, TerminatorKind}; +use stable_mir::ty::{FnDef, RigidTy, TyKind}; + +const CRATE_NAME: &str = "crate_closure_body"; + +fn test_closure_body() -> ControlFlow<()> { + let crate_items = stable_mir::all_local_items(); + for item in crate_items { + let item_ty = item.ty(); + match &item_ty.kind() { + TyKind::RigidTy(RigidTy::Closure(closure_def, _)) => { + let closure_body = closure_def.body().unwrap(); + check_incr_closure_body(closure_body); + } + _ => {} + } + } + + ControlFlow::Continue(()) +} + +fn check_incr_closure_body(body: Body) { + let first_block = &body.blocks[0]; + let TerminatorKind::Call { func: Operand::Constant(ConstOperand { const_, .. }), .. } = + &first_block.terminator.kind + else { + panic!("expected Call Terminator, got: "); + }; + + let TyKind::RigidTy(RigidTy::FnDef(FnDef(def_id), ..), ..) = const_.ty().kind() else { + panic!("expected FnDef"); + }; + + assert_eq!(def_id.name(), "incr"); +} + +fn main() { + let path = "closure_body.rs"; + generate_input(&path).unwrap(); + let args = &[ + "rustc".to_string(), + "-Cpanic=abort".to_string(), + "--crate-name".to_string(), + CRATE_NAME.to_string(), + path.to_string(), + ]; + run!(args, test_closure_body).unwrap(); +} + +fn generate_input(path: &str) -> std::io::Result<()> { + let mut file = std::fs::File::create(path)?; + write!( + file, + r#" + fn incr(y: i32) -> i32 {{ + y + 1 + }} + + fn main() {{ + let cl_incr = |x: i32| {{ + incr(x) + }}; + + let _= cl_incr(5); + }} + "# + )?; + Ok(()) +}