From fa57a48cf589fd01f7cd9721f13281af8292168f Mon Sep 17 00:00:00 2001
From: ouz-a <ouz.agz@gmail.com>
Date: Tue, 12 Sep 2023 11:50:50 +0300
Subject: [PATCH] span is index

---
 compiler/rustc_smir/src/rustc_internal/mod.rs | 17 ++++++++++++++++-
 compiler/rustc_smir/src/rustc_smir/mod.rs     | 12 +++++++-----
 compiler/rustc_smir/src/stable_mir/mod.rs     |  2 +-
 compiler/rustc_smir/src/stable_mir/ty.rs      | 11 ++++++++++-
 4 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index 4203990ab0560..10ee5af86c627 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -17,6 +17,7 @@ use rustc_interface::{interface, Queries};
 use rustc_middle::mir::interpret::AllocId;
 use rustc_middle::ty::TyCtxt;
 pub use rustc_span::def_id::{CrateNum, DefId};
+use rustc_span::Span;
 
 fn with_tables<R>(mut f: impl FnMut(&mut Tables<'_>) -> R) -> R {
     let mut ret = None;
@@ -159,6 +160,17 @@ impl<'tcx> Tables<'tcx> {
         self.alloc_ids.push(aid);
         stable_mir::AllocId(id)
     }
+
+    pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
+        for (i, &sp) in self.spans.iter().enumerate() {
+            if sp == span {
+                return stable_mir::ty::Span(i);
+            }
+        }
+        let id = self.spans.len();
+        self.spans.push(span);
+        stable_mir::ty::Span(id)
+    }
 }
 
 pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
@@ -166,7 +178,10 @@ pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
 }
 
 pub fn run(tcx: TyCtxt<'_>, f: impl FnOnce()) {
-    crate::stable_mir::run(Tables { tcx, def_ids: vec![], alloc_ids: vec![], types: vec![] }, f);
+    crate::stable_mir::run(
+        Tables { tcx, def_ids: vec![], alloc_ids: vec![], spans: vec![], types: vec![] },
+        f,
+    );
 }
 
 /// A type that provides internal information but that can still be used for debug purpose.
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 321ee7905093e..c5c32cd6e2bdf 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -9,7 +9,9 @@
 
 use crate::rustc_internal::{self, opaque};
 use crate::stable_mir::mir::{CopyNonOverlapping, UserTypeProjection, VariantIdx};
-use crate::stable_mir::ty::{FloatTy, GenericParamDef, IntTy, Movability, RigidTy, TyKind, UintTy};
+use crate::stable_mir::ty::{
+    FloatTy, GenericParamDef, IntTy, Movability, RigidTy, Span, TyKind, UintTy,
+};
 use crate::stable_mir::{self, CompilerError, Context};
 use rustc_hir as hir;
 use rustc_middle::mir::interpret::{alloc_range, AllocId};
@@ -42,7 +44,7 @@ impl<'tcx> Context for Tables<'tcx> {
         self.tcx.def_path_str(self[def_id])
     }
 
-    fn span_of_an_item(&mut self, def_id: stable_mir::DefId) -> stable_mir::ty::Span {
+    fn span_of_an_item(&mut self, def_id: stable_mir::DefId) -> Span {
         self.tcx.def_span(self[def_id]).stable(self)
     }
 
@@ -168,6 +170,7 @@ pub struct Tables<'tcx> {
     pub tcx: TyCtxt<'tcx>,
     pub def_ids: Vec<DefId>,
     pub alloc_ids: Vec<AllocId>,
+    pub spans: Vec<rustc_span::Span>,
     pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,
 }
 
@@ -1497,9 +1500,8 @@ impl<'tcx> Stable<'tcx> for ty::Region<'tcx> {
 impl<'tcx> Stable<'tcx> for rustc_span::Span {
     type T = stable_mir::ty::Span;
 
-    fn stable(&self, _: &mut Tables<'tcx>) -> Self::T {
-        // FIXME: add a real implementation of stable spans
-        opaque(self)
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.create_span(*self)
     }
 }
 
diff --git a/compiler/rustc_smir/src/stable_mir/mod.rs b/compiler/rustc_smir/src/stable_mir/mod.rs
index 508e3aa129136..09db44614e89f 100644
--- a/compiler/rustc_smir/src/stable_mir/mod.rs
+++ b/compiler/rustc_smir/src/stable_mir/mod.rs
@@ -90,7 +90,7 @@ impl CrateItem {
         with(|cx| cx.mir_body(self.0))
     }
 
-    pub fn span(&self) -> ty::Span {
+    pub fn span(&self) -> Span {
         with(|cx| cx.span_of_an_item(self.0))
     }
 }
diff --git a/compiler/rustc_smir/src/stable_mir/ty.rs b/compiler/rustc_smir/src/stable_mir/ty.rs
index 54cba1263b7b1..5a63e13f8ab92 100644
--- a/compiler/rustc_smir/src/stable_mir/ty.rs
+++ b/compiler/rustc_smir/src/stable_mir/ty.rs
@@ -35,7 +35,16 @@ pub struct Const {
 
 type Ident = Opaque;
 pub(crate) type Region = Opaque;
-pub(crate) type Span = Opaque;
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub struct Span(pub(crate) usize);
+
+impl Debug for Span {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        let mut span = None;
+        with(|context| context.rustc_tables(&mut |tables| span = Some(tables.spans[self.0])));
+        f.write_fmt(format_args!("{:?}", &span.unwrap()))
+    }
+}
 
 #[derive(Clone, Debug)]
 pub enum TyKind {