Allow fetching the SMIR body of FnDefs

3 files changed, 15 insertions, 5 deletions

diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index c50cc7cd7cd..a7640736481 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -78,8 +78,8 @@ impl<'tcx> Context for Tables<'tcx> {
         impl_trait.stable(self)
     }
 
-    fn mir_body(&mut self, item: &stable_mir::CrateItem) -> stable_mir::mir::Body {
-        let def_id = self[item.0];
+    fn mir_body(&mut self, item: stable_mir::DefId) -> stable_mir::mir::Body {
+        let def_id = self[item];
         let mir = self.tcx.optimized_mir(def_id);
         stable_mir::mir::Body {
             blocks: mir
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 90fc569f47b..9cab7230b97 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -87,7 +87,7 @@ pub struct CrateItem(pub(crate) DefId);
 
 impl CrateItem {
     pub fn body(&self) -> mir::Body {
-        with(|cx| cx.mir_body(self))
+        with(|cx| cx.mir_body(self.0))
     }
 }
 
@@ -138,7 +138,7 @@ pub trait Context {
     fn entry_fn(&mut self) -> Option<CrateItem>;
     /// Retrieve all items of the local crate that have a MIR associated with them.
     fn all_local_items(&mut self) -> CrateItems;
-    fn mir_body(&mut self, item: &CrateItem) -> mir::Body;
+    fn mir_body(&mut self, item: DefId) -> mir::Body;
     fn all_trait_decls(&mut self) -> TraitDecls;
     fn trait_decl(&mut self, trait_def: &TraitDef) -> TraitDecl;
     fn all_trait_impls(&mut self) -> ImplTraitDecls;
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index 7c58b640ec7..ad7a67275ee 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -1,4 +1,8 @@
-use super::{mir::Mutability, mir::Safety, with, AllocId, DefId};
+use super::{
+    mir::Safety,
+    mir::{Body, Mutability},
+    with, AllocId, DefId,
+};
 use crate::rustc_internal::Opaque;
 
 #[derive(Copy, Clone, Debug)]
@@ -95,6 +99,12 @@ pub struct ForeignDef(pub(crate) DefId);
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub struct FnDef(pub(crate) DefId);
 
+impl FnDef {
+    pub fn body(&self) -> Body {
+        with(|ctx| ctx.mir_body(self.0))
+    }
+}
+
 #[derive(Clone, Copy, PartialEq, Eq, Debug)]
 pub struct ClosureDef(pub(crate) DefId);