From a40dd109dc85697b5b7fcbc0dcbd0a17083568fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?O=C4=9Fuz=20A=C4=9Fcayaz=C4=B1?= <ouz.agz@gmail.com>
Date: Wed, 11 Oct 2023 12:44:59 +0300
Subject: [PATCH] implement get_filename/lines for span

---
 compiler/rustc_smir/src/rustc_smir/mod.rs | 32 ++++++++++++++++++++++-
 compiler/stable_mir/src/lib.rs            | 17 ++++++++----
 compiler/stable_mir/src/ty.rs             | 21 ++++++++++++++-
 3 files changed, 63 insertions(+), 7 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 2a265fc1f5bc0..6876b23b35ca5 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -18,7 +18,7 @@ use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use rustc_target::abi::FieldIdx;
 use stable_mir::mir::{CopyNonOverlapping, Statement, UserTypeProjection, VariantIdx};
 use stable_mir::ty::{FloatTy, GenericParamDef, IntTy, Movability, RigidTy, Span, TyKind, UintTy};
-use stable_mir::{self, opaque, Context};
+use stable_mir::{self, opaque, Context, Filename};
 use tracing::debug;
 
 mod alloc;
@@ -54,6 +54,36 @@ impl<'tcx> Context for Tables<'tcx> {
         self.tcx.sess.source_map().span_to_diagnostic_string(self[span])
     }
 
+    fn get_filename(&self, span: &Span) -> Filename {
+        opaque(
+            &self
+                .tcx
+                .sess
+                .source_map()
+                .span_to_filename(self[*span])
+                .display(rustc_span::FileNameDisplayPreference::Short)
+                .to_string(),
+        )
+    }
+
+    fn get_lines(&self, span: &Span) -> Vec<stable_mir::ty::LineInfo> {
+        let lines = &self
+            .tcx
+            .sess
+            .source_map()
+            .span_to_lines(self[*span])
+            .unwrap()
+            .lines
+            .iter()
+            .map(|line| stable_mir::ty::LineInfo {
+                line_index: line.line_index,
+                start_col: line.start_col.0,
+                end_col: line.end_col.0,
+            })
+            .collect::<Vec<stable_mir::ty::LineInfo>>();
+        lines.to_vec()
+    }
+
     fn def_kind(&mut self, def_id: stable_mir::DefId) -> stable_mir::DefKind {
         self.tcx.def_kind(self[def_id]).stable(self)
     }
diff --git a/compiler/stable_mir/src/lib.rs b/compiler/stable_mir/src/lib.rs
index f371f46204fe4..8ff2f16327a87 100644
--- a/compiler/stable_mir/src/lib.rs
+++ b/compiler/stable_mir/src/lib.rs
@@ -22,8 +22,8 @@ use std::fmt;
 use std::fmt::Debug;
 
 use self::ty::{
-    GenericPredicates, Generics, ImplDef, ImplTrait, IndexedVal, Span, TraitDecl, TraitDef, Ty,
-    TyKind,
+    GenericPredicates, Generics, ImplDef, ImplTrait, IndexedVal, LineInfo, Span, TraitDecl,
+    TraitDef, Ty, TyKind,
 };
 
 #[macro_use]
@@ -108,6 +108,7 @@ pub struct Crate {
 }
 
 pub type DefKind = Opaque;
+pub type Filename = Opaque;
 
 /// Holds information about an item in the crate.
 /// For now, it only stores the item DefId. Use functions inside `rustc_internal` module to
@@ -196,13 +197,19 @@ pub trait Context {
     /// Find a crate with the given name.
     fn find_crates(&self, name: &str) -> Vec<Crate>;
 
-    /// Prints the name of given `DefId`
+    /// Returns the name of given `DefId`
     fn name_of_def_id(&self, def_id: DefId) -> String;
 
-    /// Prints a human readable form of `Span`
+    /// Returns printable, human readable form of `Span`
     fn print_span(&self, span: Span) -> String;
 
-    /// Prints the kind of given `DefId`
+    /// Return filename from given `Span`, for diagnostic purposes
+    fn get_filename(&self, span: &Span) -> Filename;
+
+    /// Return lines corresponding to this `Span`
+    fn get_lines(&self, span: &Span) -> Vec<LineInfo>;
+
+    /// Returns the `kind` of given `DefId`
     fn def_kind(&mut self, def_id: DefId) -> DefKind;
 
     /// `Span` of an item
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index 691af15da8ca8..c6234061a6c5a 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -3,7 +3,7 @@ use super::{
     mir::{Body, Mutability},
     with, AllocId, DefId, Symbol,
 };
-use crate::Opaque;
+use crate::{Filename, Opaque};
 use std::fmt::{self, Debug, Formatter};
 
 #[derive(Copy, Clone)]
@@ -86,6 +86,25 @@ impl Debug for Span {
     }
 }
 
+impl Span {
+    /// Return filename for diagnostic purposes
+    pub fn get_filename(&self) -> Filename {
+        with(|c| c.get_filename(self))
+    }
+
+    /// Return lines that coresspond to this `Span`
+    pub fn get_lines(&self) -> Vec<LineInfo> {
+        with(|c| c.get_lines(&self))
+    }
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct LineInfo {
+    pub line_index: usize,
+    pub start_col: usize,
+    pub end_col: usize,
+}
+
 impl IndexedVal for Span {
     fn to_val(index: usize) -> Self {
         Span(index)