From 36bb79fc5e7b0c6e50d631acd4ed575a7e10da0a Mon Sep 17 00:00:00 2001
From: "Celina G. Val" <celinval@amazon.com>
Date: Mon, 18 Dec 2023 17:10:16 -0800
Subject: [PATCH] Add the function body span to StableMIR

---
 compiler/rustc_smir/src/rustc_smir/convert/mir.rs | 1 +
 compiler/stable_mir/src/mir/body.rs               | 6 +++++-
 compiler/stable_mir/src/mir/visit.rs              | 4 +++-
 3 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_smir/convert/mir.rs b/compiler/rustc_smir/src/rustc_smir/convert/mir.rs
index 41ab4007a67..49bf2192f82 100644
--- a/compiler/rustc_smir/src/rustc_smir/convert/mir.rs
+++ b/compiler/rustc_smir/src/rustc_smir/convert/mir.rs
@@ -37,6 +37,7 @@ fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
             self.arg_count,
             self.var_debug_info.iter().map(|info| info.stable(tables)).collect(),
             self.spread_arg.stable(tables),
+            self.span.stable(tables),
         )
     }
 }
diff --git a/compiler/stable_mir/src/mir/body.rs b/compiler/stable_mir/src/mir/body.rs
index 5023af9ab79..b8fd9370aa6 100644
--- a/compiler/stable_mir/src/mir/body.rs
+++ b/compiler/stable_mir/src/mir/body.rs
@@ -27,6 +27,9 @@ pub struct Body {
     ///
     /// This is used for the "rust-call" ABI such as closures.
     pub(super) spread_arg: Option<Local>,
+
+    /// The span that covers the entire function body.
+    pub span: Span,
 }
 
 pub type BasicBlockIdx = usize;
@@ -42,6 +45,7 @@ pub fn new(
         arg_count: usize,
         var_debug_info: Vec<VarDebugInfo>,
         spread_arg: Option<Local>,
+        span: Span,
     ) -> Self {
         // If locals doesn't contain enough entries, it can lead to panics in
         // `ret_local`, `arg_locals`, and `inner_locals`.
@@ -49,7 +53,7 @@ pub fn new(
             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, var_debug_info, spread_arg }
+        Self { blocks, locals, arg_count, var_debug_info, spread_arg, span }
     }
 
     /// Return local that holds this function's return value.
diff --git a/compiler/stable_mir/src/mir/visit.rs b/compiler/stable_mir/src/mir/visit.rs
index 98336a72900..ab57ff0f8f5 100644
--- a/compiler/stable_mir/src/mir/visit.rs
+++ b/compiler/stable_mir/src/mir/visit.rs
@@ -133,7 +133,7 @@ fn visit_var_debug_info(&mut self, var_debug_info: &VarDebugInfo) {
     }
 
     fn super_body(&mut self, body: &Body) {
-        let Body { blocks, locals: _, arg_count, var_debug_info, spread_arg: _ } = body;
+        let Body { blocks, locals: _, arg_count, var_debug_info, spread_arg: _, span } = body;
 
         for bb in blocks {
             self.visit_basic_block(bb);
@@ -153,6 +153,8 @@ fn super_body(&mut self, body: &Body) {
         for info in var_debug_info.iter() {
             self.visit_var_debug_info(info);
         }
+
+        self.visit_span(span)
     }
 
     fn super_basic_block(&mut self, bb: &BasicBlock) {