From e4c41b07f0b7fd15e19818b3349bb4bf726342d2 Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Mon, 23 Oct 2023 18:07:07 +0100 Subject: [PATCH 1/9] Add arg_count field to Body in Stable MIR This field allows SMIR consumers to identify which locals correspond to argument locals. It simply exposes the arg_count field from the MIR representation. --- compiler/rustc_smir/src/rustc_smir/mod.rs | 1 + compiler/stable_mir/src/mir/body.rs | 10 ++++++++++ 2 files changed, 11 insertions(+) diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs index e772ae942fa5a..86f3c7c2e3fc6 100644 --- a/compiler/rustc_smir/src/rustc_smir/mod.rs +++ b/compiler/rustc_smir/src/rustc_smir/mod.rs @@ -308,6 +308,7 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> { span: decl.source_info.span.stable(tables), }) .collect(), + arg_count: self.arg_count, } } } diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index fc617513aeeed..34dff7a621448 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -2,10 +2,20 @@ use crate::ty::{AdtDef, ClosureDef, Const, CoroutineDef, GenericArgs, Movability use crate::Opaque; use crate::{ty::Ty, Span}; +/// The SMIR representation of a single function. #[derive(Clone, Debug)] pub struct Body { pub blocks: Vec, + + /// Declarations of locals. + /// + /// The first local is the return value pointer, followed by `arg_count` + /// locals for the function arguments, followed by any user-declared + /// variables and temporaries. pub locals: LocalDecls, + + /// The number of arguments this function takes. + pub arg_count: usize, } type LocalDecls = Vec; From 93d1b3e92a1cd19e1609a4755a68b62e8f528707 Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Wed, 25 Oct 2023 10:21:47 +0100 Subject: [PATCH 2/9] Replace arg_count in public API with return/arg getters This commit hides the arg_count field in Body and instead exposes more stable and user-friendly methods to get the return and argument locals. As a result, Body instances must now be constructed using the `new` function. --- compiler/rustc_smir/src/rustc_smir/mod.rs | 12 +++---- compiler/stable_mir/src/mir/body.rs | 38 +++++++++++++++++++---- 2 files changed, 37 insertions(+), 13 deletions(-) diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs index 86f3c7c2e3fc6..1eb8b0ca40617 100644 --- a/compiler/rustc_smir/src/rustc_smir/mod.rs +++ b/compiler/rustc_smir/src/rustc_smir/mod.rs @@ -287,9 +287,8 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> { type T = stable_mir::mir::Body; fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T { - stable_mir::mir::Body { - blocks: self - .basic_blocks + stable_mir::mir::Body::new( + self.basic_blocks .iter() .map(|block| stable_mir::mir::BasicBlock { terminator: block.terminator().stable(tables), @@ -300,16 +299,15 @@ impl<'tcx> Stable<'tcx> for mir::Body<'tcx> { .collect(), }) .collect(), - locals: self - .local_decls + self.local_decls .iter() .map(|decl| stable_mir::mir::LocalDecl { ty: decl.ty.stable(tables), span: decl.source_info.span.stable(tables), }) .collect(), - arg_count: self.arg_count, - } + self.arg_count, + ) } } diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 34dff7a621448..8cdc9614db672 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -8,14 +8,40 @@ pub struct Body { pub blocks: Vec, /// Declarations of locals. - /// - /// The first local is the return value pointer, followed by `arg_count` - /// locals for the function arguments, followed by any user-declared - /// variables and temporaries. + // + // The first local is the return value pointer, followed by `arg_count` + // locals for the function arguments, followed by any user-declared + // variables and temporaries. pub locals: LocalDecls, - /// The number of arguments this function takes. - pub arg_count: usize, + // The number of arguments this function takes. + arg_count: usize, +} + +impl Body { + /// Constructs a `Body`. + /// + /// A constructor is required to build a `Body` from outside the crate + /// because the `arg_count` field is private. + pub fn new(blocks: Vec, locals: LocalDecls, arg_count: usize) -> Self { + // If locals doesn't contain enough entries, it can lead to panics in + // `ret_local` and `arg_locals`. + assert!( + locals.len() >= arg_count + 1, + "A Body must contain at least a local for the return value and each of the function's arguments" + ); + Self { blocks, locals, arg_count } + } + + /// Gets the function's return local. + pub fn ret_local(&self) -> &LocalDecl { + &self.locals[0] + } + + /// Gets the locals in `self` that correspond to the function's arguments. + pub fn arg_locals(&self) -> &[LocalDecl] { + &self.locals[1..self.arg_count + 1] + } } type LocalDecls = Vec; From f4d80a5f09afcbe0bd80147a0685be1dfaf81acd Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Wed, 25 Oct 2023 16:51:18 +0100 Subject: [PATCH 3/9] Add public API to retrieve internal locals --- compiler/stable_mir/src/mir/body.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 8cdc9614db672..f76ad87584f58 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -33,15 +33,21 @@ impl Body { Self { blocks, locals, arg_count } } - /// Gets the function's return local. + /// Return local that holds this function's return value. pub fn ret_local(&self) -> &LocalDecl { &self.locals[0] } - /// Gets the locals in `self` that correspond to the function's arguments. + /// Locals in `self` that correspond to this function's arguments. pub fn arg_locals(&self) -> &[LocalDecl] { &self.locals[1..self.arg_count + 1] } + + /// Internal locals for this function. These are the locals that are + /// neither the return local nor the argument locals. + pub fn internal_locals(&self) -> &[LocalDecl] { + &self.locals[self.arg_count + 1..] + } } type LocalDecls = Vec; From 372c5335327f330858b03889a46614d2a248d325 Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Wed, 25 Oct 2023 21:28:18 +0100 Subject: [PATCH 4/9] Make locals field private --- compiler/stable_mir/src/mir/body.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index f76ad87584f58..dc472b3786c85 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -7,12 +7,12 @@ use crate::{ty::Ty, Span}; pub struct Body { pub blocks: Vec, - /// Declarations of locals. + // Declarations of locals within the function. // // The first local is the return value pointer, followed by `arg_count` // locals for the function arguments, followed by any user-declared // variables and temporaries. - pub locals: LocalDecls, + locals: LocalDecls, // The number of arguments this function takes. arg_count: usize, @@ -22,12 +22,12 @@ impl Body { /// Constructs a `Body`. /// /// A constructor is required to build a `Body` from outside the crate - /// because the `arg_count` field is private. + /// because the `arg_count` and `locals` fields are private. pub fn new(blocks: Vec, locals: LocalDecls, arg_count: usize) -> Self { // If locals doesn't contain enough entries, it can lead to panics in - // `ret_local` and `arg_locals`. + // `ret_local`, `arg_locals`, and `internal_locals`. assert!( - locals.len() >= arg_count + 1, + locals.len() > arg_count, "A Body must contain at least a local for the return value and each of the function's arguments" ); Self { blocks, locals, arg_count } From 39b293fb5a68081e7c050ed2805c8e3404c9763a Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Wed, 25 Oct 2023 22:09:12 +0100 Subject: [PATCH 5/9] Add a public API to get all body locals This is particularly helpful for the ui tests, but also could be helpful for Stable MIR users who just want all the locals without needing to concatenate responses --- compiler/stable_mir/src/mir/body.rs | 8 ++++++++ tests/ui-fulldeps/stable-mir/check_instance.rs | 2 +- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index dc472b3786c85..97dd649bf64ab 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -48,6 +48,14 @@ impl Body { pub fn internal_locals(&self) -> &[LocalDecl] { &self.locals[self.arg_count + 1..] } + + /// Convenience function to get all the locals in this function. + /// + /// Locals are typically accessed via the more specific methods `ret_local`, + /// `arg_locals`, and `internal_locals`. + pub fn locals(&self) -> &[LocalDecl] { + &self.locals + } } type LocalDecls = Vec; diff --git a/tests/ui-fulldeps/stable-mir/check_instance.rs b/tests/ui-fulldeps/stable-mir/check_instance.rs index ee82bc77aedae..29d88b31e77b3 100644 --- a/tests/ui-fulldeps/stable-mir/check_instance.rs +++ b/tests/ui-fulldeps/stable-mir/check_instance.rs @@ -59,7 +59,7 @@ fn test_body(body: mir::Body) { for term in body.blocks.iter().map(|bb| &bb.terminator) { match &term.kind { Call { func, .. } => { - let TyKind::RigidTy(ty) = func.ty(&body.locals).kind() else { unreachable!() }; + let TyKind::RigidTy(ty) = func.ty(&body.locals()).kind() else { unreachable!() }; let RigidTy::FnDef(def, args) = ty else { unreachable!() }; let result = Instance::resolve(def, &args); assert!(result.is_ok()); From fe4dfb814b68bb03468736c5c0632f6495ea855e Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Thu, 26 Oct 2023 00:18:42 +0100 Subject: [PATCH 6/9] Rename internal_locals to inner_locals The word internal has connotations about information that's not exposed. It's more accurate to say that the remaining locals apply only to the inner part of the function, so I'm renaming them to inner locals. --- compiler/stable_mir/src/mir/body.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 97dd649bf64ab..75c988056b445 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -25,7 +25,7 @@ impl Body { /// because the `arg_count` and `locals` fields are private. pub fn new(blocks: Vec, locals: LocalDecls, arg_count: usize) -> Self { // If locals doesn't contain enough entries, it can lead to panics in - // `ret_local`, `arg_locals`, and `internal_locals`. + // `ret_local`, `arg_locals`, and `inner_locals`. assert!( locals.len() > arg_count, "A Body must contain at least a local for the return value and each of the function's arguments" @@ -43,16 +43,16 @@ impl Body { &self.locals[1..self.arg_count + 1] } - /// Internal locals for this function. These are the locals that are + /// Inner locals for this function. These are the locals that are /// neither the return local nor the argument locals. - pub fn internal_locals(&self) -> &[LocalDecl] { + pub fn inner_locals(&self) -> &[LocalDecl] { &self.locals[self.arg_count + 1..] } /// Convenience function to get all the locals in this function. /// /// Locals are typically accessed via the more specific methods `ret_local`, - /// `arg_locals`, and `internal_locals`. + /// `arg_locals`, and `inner_locals`. pub fn locals(&self) -> &[LocalDecl] { &self.locals } From 4b23bd47348a6325ee332e48b9d16e1b65d1c500 Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Thu, 26 Oct 2023 00:21:28 +0100 Subject: [PATCH 7/9] Update Place and Operand to take slices The latest locals() method in stable MIR returns slices instead of vecs. This commit also includes fixes to the existing tests that previously referenced the private locals field. --- compiler/stable_mir/src/mir/body.rs | 4 ++-- .../ui-fulldeps/stable-mir/check_instance.rs | 2 +- tests/ui-fulldeps/stable-mir/crate-info.rs | 22 +++++++++---------- tests/ui-fulldeps/stable-mir/smir_internal.rs | 2 +- 4 files changed, 15 insertions(+), 15 deletions(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 75c988056b445..8081def147924 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -517,7 +517,7 @@ pub enum NullOp { } impl Operand { - pub fn ty(&self, locals: &LocalDecls) -> Ty { + pub fn ty(&self, locals: &[LocalDecl]) -> Ty { match self { Operand::Copy(place) | Operand::Move(place) => place.ty(locals), Operand::Constant(c) => c.ty(), @@ -532,7 +532,7 @@ impl Constant { } impl Place { - pub fn ty(&self, locals: &LocalDecls) -> Ty { + pub fn ty(&self, locals: &[LocalDecl]) -> Ty { let _start_ty = locals[self.local].ty; todo!("Implement projection") } diff --git a/tests/ui-fulldeps/stable-mir/check_instance.rs b/tests/ui-fulldeps/stable-mir/check_instance.rs index 29d88b31e77b3..a340877752d8f 100644 --- a/tests/ui-fulldeps/stable-mir/check_instance.rs +++ b/tests/ui-fulldeps/stable-mir/check_instance.rs @@ -59,7 +59,7 @@ fn test_body(body: mir::Body) { for term in body.blocks.iter().map(|bb| &bb.terminator) { match &term.kind { Call { func, .. } => { - let TyKind::RigidTy(ty) = func.ty(&body.locals()).kind() else { unreachable!() }; + let TyKind::RigidTy(ty) = func.ty(body.locals()).kind() else { unreachable!() }; let RigidTy::FnDef(def, args) = ty else { unreachable!() }; let result = Instance::resolve(def, &args); assert!(result.is_ok()); diff --git a/tests/ui-fulldeps/stable-mir/crate-info.rs b/tests/ui-fulldeps/stable-mir/crate-info.rs index 60c6053d295b4..33673cd131ff3 100644 --- a/tests/ui-fulldeps/stable-mir/crate-info.rs +++ b/tests/ui-fulldeps/stable-mir/crate-info.rs @@ -47,7 +47,7 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> { let bar = get_item(&items, (DefKind::Fn, "bar")).unwrap(); let body = bar.body(); - assert_eq!(body.locals.len(), 2); + assert_eq!(body.locals().len(), 2); assert_eq!(body.blocks.len(), 1); let block = &body.blocks[0]; assert_eq!(block.statements.len(), 1); @@ -62,7 +62,7 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> { let foo_bar = get_item(&items, (DefKind::Fn, "foo_bar")).unwrap(); let body = foo_bar.body(); - assert_eq!(body.locals.len(), 5); + assert_eq!(body.locals().len(), 5); assert_eq!(body.blocks.len(), 4); let block = &body.blocks[0]; match &block.terminator.kind { @@ -72,29 +72,29 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> { let types = get_item(&items, (DefKind::Fn, "types")).unwrap(); let body = types.body(); - assert_eq!(body.locals.len(), 6); + assert_eq!(body.locals().len(), 6); assert_matches!( - body.locals[0].ty.kind(), + body.locals()[0].ty.kind(), stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Bool) ); assert_matches!( - body.locals[1].ty.kind(), + body.locals()[1].ty.kind(), stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Bool) ); assert_matches!( - body.locals[2].ty.kind(), + body.locals()[2].ty.kind(), stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Char) ); assert_matches!( - body.locals[3].ty.kind(), + body.locals()[3].ty.kind(), stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Int(stable_mir::ty::IntTy::I32)) ); assert_matches!( - body.locals[4].ty.kind(), + body.locals()[4].ty.kind(), stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Uint(stable_mir::ty::UintTy::U64)) ); assert_matches!( - body.locals[5].ty.kind(), + body.locals()[5].ty.kind(), stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Float( stable_mir::ty::FloatTy::F64 )) @@ -123,10 +123,10 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> { for block in instance.body().blocks { match &block.terminator.kind { stable_mir::mir::TerminatorKind::Call { func, .. } => { - let TyKind::RigidTy(ty) = func.ty(&body.locals).kind() else { unreachable!() }; + let TyKind::RigidTy(ty) = func.ty(&body.locals()).kind() else { unreachable!() }; let RigidTy::FnDef(def, args) = ty else { unreachable!() }; let next_func = Instance::resolve(def, &args).unwrap(); - match next_func.body().locals[1].ty.kind() { + match next_func.body().locals()[1].ty.kind() { TyKind::RigidTy(RigidTy::Uint(_)) | TyKind::RigidTy(RigidTy::Tuple(_)) => {} other => panic!("{other:?}"), } diff --git a/tests/ui-fulldeps/stable-mir/smir_internal.rs b/tests/ui-fulldeps/stable-mir/smir_internal.rs index 5ad05559cb4bb..b0596b1882383 100644 --- a/tests/ui-fulldeps/stable-mir/smir_internal.rs +++ b/tests/ui-fulldeps/stable-mir/smir_internal.rs @@ -29,7 +29,7 @@ const CRATE_NAME: &str = "input"; fn test_translation(_tcx: TyCtxt<'_>) -> ControlFlow<()> { let main_fn = stable_mir::entry_fn().unwrap(); let body = main_fn.body(); - let orig_ty = body.locals[0].ty; + let orig_ty = body.locals()[0].ty; let rustc_ty = rustc_internal::internal(&orig_ty); assert!(rustc_ty.is_unit()); ControlFlow::Continue(()) From bac7d5b52cc60754e0555f96c3501e1531bbf6fe Mon Sep 17 00:00:00 2001 From: Kirby Linvill Date: Thu, 26 Oct 2023 00:22:56 +0100 Subject: [PATCH 8/9] Add test for smir locals --- tests/ui-fulldeps/stable-mir/crate-info.rs | 31 ++++++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/tests/ui-fulldeps/stable-mir/crate-info.rs b/tests/ui-fulldeps/stable-mir/crate-info.rs index 33673cd131ff3..ed6b786f5e1de 100644 --- a/tests/ui-fulldeps/stable-mir/crate-info.rs +++ b/tests/ui-fulldeps/stable-mir/crate-info.rs @@ -140,6 +140,29 @@ fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> { // Ensure we don't panic trying to get the body of a constant. foo_const.body(); + let locals_fn = get_item(&items, (DefKind::Fn, "locals")).unwrap(); + let body = locals_fn.body(); + assert_eq!(body.locals().len(), 4); + assert_matches!( + body.ret_local().ty.kind(), + stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Char) + ); + assert_eq!(body.arg_locals().len(), 2); + assert_matches!( + body.arg_locals()[0].ty.kind(), + stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Int(stable_mir::ty::IntTy::I32)) + ); + assert_matches!( + body.arg_locals()[1].ty.kind(), + stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Uint(stable_mir::ty::UintTy::U64)) + ); + assert_eq!(body.inner_locals().len(), 1); + // If conditions have an extra inner local to hold their results + assert_matches!( + body.inner_locals()[0].ty.kind(), + stable_mir::ty::TyKind::RigidTy(stable_mir::ty::RigidTy::Bool) + ); + ControlFlow::Continue(()) } @@ -211,6 +234,14 @@ fn generate_input(path: &str) -> std::io::Result<()> { pub fn assert(x: i32) -> i32 {{ x + 1 + }} + + pub fn locals(a: i32, _: u64) -> char {{ + if a > 5 {{ + 'a' + }} else {{ + 'b' + }} }}"# )?; Ok(()) From d55487d7e9834c8fe01350b5b4fa402423c855d5 Mon Sep 17 00:00:00 2001 From: Oli Scherer Date: Thu, 26 Oct 2023 12:32:47 +0200 Subject: [PATCH 9/9] Use two slice expressions to save on an offset repetition --- compiler/stable_mir/src/mir/body.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs index 8081def147924..d86d56b3dc724 100644 --- a/compiler/stable_mir/src/mir/body.rs +++ b/compiler/stable_mir/src/mir/body.rs @@ -40,7 +40,7 @@ impl Body { /// Locals in `self` that correspond to this function's arguments. pub fn arg_locals(&self) -> &[LocalDecl] { - &self.locals[1..self.arg_count + 1] + &self.locals[1..][..self.arg_count] } /// Inner locals for this function. These are the locals that are