Rollup merge of #118012 - celinval:smir-alloc, r=ouz-a

Add support for global allocation in smir

Add APIs to StableMir to support global allocation. Before this change, StableMir users had no API available to retrieve Allocation provenance information. They had to resource to internal APIs instead.

One example is retrieving the Allocation of an `&str`. See test for an example on how the API can be used.

9 files changed, 329 insertions, 26 deletions

diff --git a/compiler/rustc_smir/src/rustc_internal/internal.rs b/compiler/rustc_smir/src/rustc_internal/internal.rs
index 24b1a3c63be..fe226ef60ed 100644
--- a/compiler/rustc_smir/src/rustc_internal/internal.rs
+++ b/compiler/rustc_smir/src/rustc_internal/internal.rs
@@ -7,12 +7,14 @@
 use crate::rustc_smir::Tables;
 use rustc_middle::ty::{self as rustc_ty, Ty as InternalTy};
 use rustc_span::Symbol;
+use stable_mir::mir::alloc::AllocId;
 use stable_mir::mir::mono::{Instance, MonoItem, StaticDef};
 use stable_mir::ty::{
-    AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const, FloatTy,
-    GenericArgKind, GenericArgs, IntTy, Region, RigidTy, TraitRef, Ty, UintTy,
+    AdtDef, Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const,
+    ExistentialTraitRef, FloatTy, GenericArgKind, GenericArgs, IntTy, Region, RigidTy, TraitRef,
+    Ty, UintTy,
 };
-use stable_mir::{AllocId, CrateItem, DefId};
+use stable_mir::{CrateItem, DefId};
 
 use super::RustcInternal;
 
@@ -228,6 +230,17 @@ impl<'tcx> RustcInternal<'tcx> for BoundVariableKind {
     }
 }
 
+impl<'tcx> RustcInternal<'tcx> for ExistentialTraitRef {
+    type T = rustc_ty::ExistentialTraitRef<'tcx>;
+
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        rustc_ty::ExistentialTraitRef {
+            def_id: self.def_id.0.internal(tables),
+            args: self.generic_args.internal(tables),
+        }
+    }
+}
+
 impl<'tcx> RustcInternal<'tcx> for TraitRef {
     type T = rustc_ty::TraitRef<'tcx>;
 
@@ -276,3 +289,13 @@ where
         (*self).internal(tables)
     }
 }
+impl<'tcx, T> RustcInternal<'tcx> for Option<T>
+where
+    T: RustcInternal<'tcx>,
+{
+    type T = Option<T::T>;
+
+    fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        self.as_ref().map(|inner| inner.internal(tables))
+    }
+}
diff --git a/compiler/rustc_smir/src/rustc_internal/mod.rs b/compiler/rustc_smir/src/rustc_internal/mod.rs
index fa75fd3076c..fed21f16a4e 100644
--- a/compiler/rustc_smir/src/rustc_internal/mod.rs
+++ b/compiler/rustc_smir/src/rustc_internal/mod.rs
@@ -118,7 +118,7 @@ impl<'tcx> Tables<'tcx> {
         self.def_ids.create_or_fetch(did)
     }
 
-    fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::AllocId {
+    pub(crate) fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::mir::alloc::AllocId {
         self.alloc_ids.create_or_fetch(aid)
     }
 
diff --git a/compiler/rustc_smir/src/rustc_smir/mod.rs b/compiler/rustc_smir/src/rustc_smir/mod.rs
index 2a9f0b42678..403c7703deb 100644
--- a/compiler/rustc_smir/src/rustc_smir/mod.rs
+++ b/compiler/rustc_smir/src/rustc_smir/mod.rs
@@ -17,15 +17,16 @@ use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::ty::{self, Instance, ParamEnv, ScalarInt, Ty, TyCtxt, Variance};
 use rustc_span::def_id::{CrateNum, DefId, LOCAL_CRATE};
 use rustc_target::abi::FieldIdx;
-use stable_mir::mir::mono::InstanceDef;
+use stable_mir::mir::alloc::GlobalAlloc;
+use stable_mir::mir::mono::{InstanceDef, StaticDef};
 use stable_mir::mir::{
     Body, ConstOperand, CopyNonOverlapping, Statement, UserTypeProjection, VarDebugInfoFragment,
     VariantIdx,
 };
 use stable_mir::ty::{
-    AdtDef, AdtKind, ClosureDef, ClosureKind, Const, ConstId, ConstantKind, EarlyParamRegion,
-    FloatTy, FnDef, GenericArgs, GenericParamDef, IntTy, LineInfo, Movability, RigidTy, Span,
-    TyKind, UintTy,
+    AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, ConstId, ConstantKind,
+    EarlyParamRegion, FloatTy, FnDef, GenericArgs, GenericParamDef, IntTy, LineInfo, Movability,
+    RigidTy, Span, TyKind, UintTy,
 };
 use stable_mir::{self, opaque, Context, CrateItem, Error, Filename, ItemKind};
 use std::cell::RefCell;
@@ -318,6 +319,30 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
             .ok_or_else(|| Error::new(format!("Const `{cnst:?}` cannot be encoded as u64")))
     }
 
+    fn eval_static_initializer(&self, def: StaticDef) -> Result<Allocation, Error> {
+        let mut tables = self.0.borrow_mut();
+        let def_id = def.0.internal(&mut *tables);
+        tables.tcx.eval_static_initializer(def_id).stable(&mut *tables)
+    }
+
+    fn global_alloc(&self, alloc: stable_mir::mir::alloc::AllocId) -> GlobalAlloc {
+        let mut tables = self.0.borrow_mut();
+        let alloc_id = alloc.internal(&mut *tables);
+        tables.tcx.global_alloc(alloc_id).stable(&mut *tables)
+    }
+
+    fn vtable_allocation(
+        &self,
+        global_alloc: &GlobalAlloc,
+    ) -> Option<stable_mir::mir::alloc::AllocId> {
+        let mut tables = self.0.borrow_mut();
+        let GlobalAlloc::VTable(ty, trait_ref) = global_alloc else { return None };
+        let alloc_id = tables
+            .tcx
+            .vtable_allocation((ty.internal(&mut *tables), trait_ref.internal(&mut *tables)));
+        Some(alloc_id.stable(&mut *tables))
+    }
+
     fn usize_to_const(&self, val: u64) -> Result<Const, Error> {
         let mut tables = self.0.borrow_mut();
         let ty = tables.tcx.types.usize;
@@ -342,7 +367,7 @@ pub(crate) struct TablesWrapper<'tcx>(pub(crate) RefCell<Tables<'tcx>>);
 pub struct Tables<'tcx> {
     pub(crate) tcx: TyCtxt<'tcx>,
     pub(crate) def_ids: IndexMap<DefId, stable_mir::DefId>,
-    pub(crate) alloc_ids: IndexMap<AllocId, stable_mir::AllocId>,
+    pub(crate) alloc_ids: IndexMap<AllocId, stable_mir::mir::alloc::AllocId>,
     pub(crate) spans: IndexMap<rustc_span::Span, Span>,
     pub(crate) types: IndexMap<Ty<'tcx>, stable_mir::ty::Ty>,
     pub(crate) instances: IndexMap<ty::Instance<'tcx>, InstanceDef>,
@@ -1590,6 +1615,14 @@ impl<'tcx> Stable<'tcx> for ty::BoundTy {
     }
 }
 
+impl<'tcx> Stable<'tcx> for mir::interpret::ConstAllocation<'tcx> {
+    type T = Allocation;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        self.inner().stable(tables)
+    }
+}
+
 impl<'tcx> Stable<'tcx> for mir::interpret::Allocation {
     type T = stable_mir::ty::Allocation;
 
@@ -1602,6 +1635,32 @@ impl<'tcx> Stable<'tcx> for mir::interpret::Allocation {
     }
 }
 
+impl<'tcx> Stable<'tcx> for mir::interpret::AllocId {
+    type T = stable_mir::mir::alloc::AllocId;
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        tables.create_alloc_id(*self)
+    }
+}
+
+impl<'tcx> Stable<'tcx> for mir::interpret::GlobalAlloc<'tcx> {
+    type T = GlobalAlloc;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        match self {
+            mir::interpret::GlobalAlloc::Function(instance) => {
+                GlobalAlloc::Function(instance.stable(tables))
+            }
+            mir::interpret::GlobalAlloc::VTable(ty, trait_ref) => {
+                GlobalAlloc::VTable(ty.stable(tables), trait_ref.stable(tables))
+            }
+            mir::interpret::GlobalAlloc::Static(def) => {
+                GlobalAlloc::Static(tables.static_def(*def))
+            }
+            mir::interpret::GlobalAlloc::Memory(alloc) => GlobalAlloc::Memory(alloc.stable(tables)),
+        }
+    }
+}
+
 impl<'tcx> Stable<'tcx> for ty::trait_def::TraitSpecializationKind {
     type T = stable_mir::ty::TraitSpecializationKind;
     fn stable(&self, _: &mut Tables<'tcx>) -> Self::T {
@@ -1989,6 +2048,14 @@ impl<'tcx> Stable<'tcx> for MonoItem<'tcx> {
     }
 }
 
+impl<'tcx> Stable<'tcx> for mir::interpret::ErrorHandled {
+    type T = Error;
+
+    fn stable(&self, _tables: &mut Tables<'tcx>) -> Self::T {
+        Error::new(format!("{self:?}"))
+    }
+}
+
 impl<'tcx, T> Stable<'tcx> for &T
 where
     T: Stable<'tcx>,
@@ -2010,3 +2077,18 @@ where
         self.as_ref().map(|value| value.stable(tables))
     }
 }
+
+impl<'tcx, T, E> Stable<'tcx> for Result<T, E>
+where
+    T: Stable<'tcx>,
+    E: Stable<'tcx>,
+{
+    type T = Result<T::T, E::T>;
+
+    fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
+        match self {
+            Ok(val) => Ok(val.stable(tables)),
+            Err(error) => Err(error.stable(tables)),
+        }
+    }
+}
diff --git a/compiler/stable_mir/src/lib.rs b/compiler/stable_mir/src/lib.rs
index 5eb6a8a5e54..6c1b723a8da 100644
--- a/compiler/stable_mir/src/lib.rs
+++ b/compiler/stable_mir/src/lib.rs
@@ -17,7 +17,7 @@
 //! The goal is to eventually be published on
 //! [crates.io](https://crates.io).
 
-use crate::mir::mono::InstanceDef;
+use crate::mir::mono::{InstanceDef, StaticDef};
 use crate::mir::Body;
 use std::fmt;
 use std::fmt::Debug;
@@ -37,9 +37,10 @@ pub mod mir;
 pub mod ty;
 pub mod visitor;
 
+use crate::mir::alloc::{AllocId, GlobalAlloc};
 use crate::mir::pretty::function_name;
 use crate::mir::Mutability;
-use crate::ty::{AdtDef, AdtKind, ClosureDef, ClosureKind, Const, RigidTy};
+use crate::ty::{AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, RigidTy};
 pub use error::*;
 use mir::mono::Instance;
 use ty::{FnDef, GenericArgs};
@@ -73,19 +74,6 @@ impl IndexedVal for DefId {
     }
 }
 
-/// A unique identification number for each provenance
-#[derive(Clone, Copy, PartialEq, Eq, Debug)]
-pub struct AllocId(usize);
-
-impl IndexedVal for AllocId {
-    fn to_val(index: usize) -> Self {
-        AllocId(index)
-    }
-    fn to_index(&self) -> usize {
-        self.0
-    }
-}
-
 /// A list of crate items.
 pub type CrateItems = Vec<CrateItem>;
 
@@ -141,6 +129,10 @@ impl CrateItem {
         with(|cx| cx.def_ty(self.0))
     }
 
+    pub fn is_foreign_item(&self) -> bool {
+        with(|cx| cx.is_foreign_item(*self))
+    }
+
     pub fn dump<W: io::Write>(&self, w: &mut W) -> io::Result<()> {
         writeln!(w, "{}", function_name(*self))?;
         self.body().dump(w)
@@ -190,6 +182,8 @@ pub fn trait_impl(trait_impl: &ImplDef) -> ImplTrait {
     with(|cx| cx.trait_impl(trait_impl))
 }
 
+/// This trait defines the interface between stable_mir and the Rust compiler.
+/// Do not use this directly.
 pub trait Context {
     fn entry_fn(&self) -> Option<CrateItem>;
     /// Retrieve all items of the local crate that have a MIR associated with them.
@@ -291,6 +285,15 @@ pub trait Context {
         args: &GenericArgs,
         kind: ClosureKind,
     ) -> Option<Instance>;
+
+    /// Evaluate a static's initializer.
+    fn eval_static_initializer(&self, def: StaticDef) -> Result<Allocation, Error>;
+
+    /// Retrieve global allocation for the given allocation ID.
+    fn global_alloc(&self, id: AllocId) -> GlobalAlloc;
+
+    /// Retrieve the id for the virtual table.
+    fn vtable_allocation(&self, global_alloc: &GlobalAlloc) -> Option<AllocId>;
 }
 
 // A thread local variable that stores a pointer to the tables mapping between TyCtxt
diff --git a/compiler/stable_mir/src/mir.rs b/compiler/stable_mir/src/mir.rs
index 2cbe6eb4ad1..82555461d64 100644
--- a/compiler/stable_mir/src/mir.rs
+++ b/compiler/stable_mir/src/mir.rs
@@ -1,3 +1,4 @@
+pub mod alloc;
 mod body;
 pub mod mono;
 pub mod pretty;
diff --git a/compiler/stable_mir/src/mir/alloc.rs b/compiler/stable_mir/src/mir/alloc.rs
new file mode 100644
index 00000000000..af951bcef8c
--- /dev/null
+++ b/compiler/stable_mir/src/mir/alloc.rs
@@ -0,0 +1,51 @@
+//! This module provides methods to retrieve allocation information, such as static variables.
+use crate::mir::mono::{Instance, StaticDef};
+use crate::ty::{Allocation, Binder, ExistentialTraitRef, IndexedVal, Ty};
+use crate::with;
+
+/// An allocation in the SMIR global memory can be either a function pointer,
+/// a static, or a "real" allocation with some data in it.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum GlobalAlloc {
+    /// The alloc ID is used as a function pointer.
+    Function(Instance),
+    /// This alloc ID points to a symbolic (not-reified) vtable.
+    /// The `None` trait ref is used to represent auto traits.
+    VTable(Ty, Option<Binder<ExistentialTraitRef>>),
+    /// The alloc ID points to a "lazy" static variable that did not get computed (yet).
+    /// This is also used to break the cycle in recursive statics.
+    Static(StaticDef),
+    /// The alloc ID points to memory.
+    Memory(Allocation),
+}
+
+impl From<AllocId> for GlobalAlloc {
+    fn from(value: AllocId) -> Self {
+        with(|cx| cx.global_alloc(value))
+    }
+}
+
+impl GlobalAlloc {
+    /// Retrieve the allocation id for a global allocation if it exists.
+    ///
+    /// For `[GlobalAlloc::VTable]`, this will return the allocation for the VTable of the given
+    /// type for the optional trait if the type implements the trait.
+    ///
+    /// This method will always return `None` for allocations other than `[GlobalAlloc::VTable]`.
+    pub fn vtable_allocation(&self) -> Option<AllocId> {
+        with(|cx| cx.vtable_allocation(self))
+    }
+}
+
+/// A unique identification number for each provenance
+#[derive(Clone, Copy, PartialEq, Eq, Debug)]
+pub struct AllocId(usize);
+
+impl IndexedVal for AllocId {
+    fn to_val(index: usize) -> Self {
+        AllocId(index)
+    }
+    fn to_index(&self) -> usize {
+        self.0
+    }
+}
diff --git a/compiler/stable_mir/src/mir/mono.rs b/compiler/stable_mir/src/mir/mono.rs
index 8c43f2c4b8f..9cec963cf84 100644
--- a/compiler/stable_mir/src/mir/mono.rs
+++ b/compiler/stable_mir/src/mir/mono.rs
@@ -1,5 +1,5 @@
 use crate::mir::Body;
-use crate::ty::{ClosureDef, ClosureKind, FnDef, GenericArgs, IndexedVal, Ty};
+use crate::ty::{Allocation, ClosureDef, ClosureKind, FnDef, GenericArgs, IndexedVal, Ty};
 use crate::{with, CrateItem, DefId, Error, ItemKind, Opaque};
 use std::fmt::{Debug, Formatter};
 
@@ -37,6 +37,11 @@ impl Instance {
         with(|context| context.instance_body(self.def))
     }
 
+    pub fn is_foreign_item(&self) -> bool {
+        let item = CrateItem::try_from(*self);
+        item.as_ref().map_or(false, CrateItem::is_foreign_item)
+    }
+
     /// Get the instance type with generic substitutions applied and lifetimes erased.
     pub fn ty(&self) -> Ty {
         with(|context| context.instance_ty(self.def))
@@ -128,6 +133,18 @@ impl From<Instance> for MonoItem {
     }
 }
 
+impl From<StaticDef> for MonoItem {
+    fn from(value: StaticDef) -> Self {
+        MonoItem::Static(value)
+    }
+}
+
+impl From<StaticDef> for CrateItem {
+    fn from(value: StaticDef) -> Self {
+        CrateItem(value.0)
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct InstanceDef(usize);
 
@@ -147,9 +164,15 @@ impl TryFrom<CrateItem> for StaticDef {
 }
 
 impl StaticDef {
+    /// Return the type of this static definition.
     pub fn ty(&self) -> Ty {
         with(|cx| cx.def_ty(self.0))
     }
+
+    /// Evaluate a static's initializer, returning the allocation of the initializer's memory.
+    pub fn eval_initializer(&self) -> Result<Allocation, Error> {
+        with(|cx| cx.eval_static_initializer(*self))
+    }
 }
 
 impl IndexedVal for InstanceDef {
diff --git a/compiler/stable_mir/src/ty.rs b/compiler/stable_mir/src/ty.rs
index abe6f8ec12f..010e2e7e4a6 100644
--- a/compiler/stable_mir/src/ty.rs
+++ b/compiler/stable_mir/src/ty.rs
@@ -1,8 +1,9 @@
 use super::{
     mir::Safety,
     mir::{Body, Mutability},
-    with, AllocId, DefId, Error, Symbol,
+    with, DefId, Error, Symbol,
 };
+use crate::mir::alloc::AllocId;
 use crate::{Filename, Opaque};
 use std::fmt::{self, Debug, Display, Formatter};
 
diff --git a/tests/ui-fulldeps/stable-mir/check_allocation.rs b/tests/ui-fulldeps/stable-mir/check_allocation.rs
new file mode 100644
index 00000000000..e5fb7311c0b
--- /dev/null
+++ b/tests/ui-fulldeps/stable-mir/check_allocation.rs
@@ -0,0 +1,119 @@
+// run-pass
+//! Test that users are able to use stable mir APIs to retrieve information of global allocations
+//! such as `vtable_allocation`.
+
+// ignore-stage1
+// ignore-cross-compile
+// ignore-remote
+// ignore-windows-gnu mingw has troubles with linking https://github.com/rust-lang/rust/pull/116837
+// edition: 2021
+
+#![feature(rustc_private)]
+#![feature(assert_matches)]
+#![feature(control_flow_enum)]
+#![feature(ascii_char, ascii_char_variants)]
+
+extern crate rustc_hir;
+extern crate rustc_middle;
+#[macro_use]
+extern crate rustc_smir;
+extern crate rustc_driver;
+extern crate rustc_interface;
+extern crate stable_mir;
+
+use rustc_middle::ty::TyCtxt;
+use rustc_smir::rustc_internal;
+use stable_mir::{CrateItem, CrateItems, ItemKind};
+use stable_mir::mir::alloc::GlobalAlloc;
+use stable_mir::mir::mono::StaticDef;
+use std::ascii::Char;
+use std::assert_matches::assert_matches;
+use std::io::Write;
+use std::ops::ControlFlow;
+
+const CRATE_NAME: &str = "input";
+
+/// This function uses the Stable MIR APIs to get information about the test crate.
+fn test_stable_mir(_tcx: TyCtxt<'_>) -> ControlFlow<()> {
+    // Find items in the local crate.
+    let items = stable_mir::all_local_items();
+    check_foo(*get_item(&items, (ItemKind::Static, "FOO")).unwrap());
+    check_bar(*get_item(&items, (ItemKind::Static, "BAR")).unwrap());
+    ControlFlow::Continue(())
+}
+
+/// Check the allocation data for static `FOO`.
+///
+/// ```no_run
+/// static FOO: [&str; 2] = ["hi", "there"];
+/// ```
+fn check_foo(item: CrateItem) {
+    let def = StaticDef::try_from(item).unwrap();
+    let alloc = def.eval_initializer().unwrap();
+    assert_eq!(alloc.provenance.ptrs.len(), 2);
+
+    let alloc_id_0 = alloc.provenance.ptrs[0].1.0;
+    assert_matches!(GlobalAlloc::from(alloc_id_0), GlobalAlloc::Memory(..));
+
+    let alloc_id_1 = alloc.provenance.ptrs[1].1.0;
+    assert_matches!(GlobalAlloc::from(alloc_id_1), GlobalAlloc::Memory(..));
+}
+
+/// Check the allocation data for static `BAR`.
+///
+/// ```no_run
+/// static BAR: &str = "Bar";
+/// ```
+fn check_bar(item: CrateItem) {
+    let def = StaticDef::try_from(item).unwrap();
+    let alloc = def.eval_initializer().unwrap();
+    assert_eq!(alloc.provenance.ptrs.len(), 1);
+
+    let alloc_id_0 = alloc.provenance.ptrs[0].1.0;
+    let GlobalAlloc::Memory(allocation) = GlobalAlloc::from(alloc_id_0) else { unreachable!() };
+    assert_eq!(allocation.bytes.len(), 3);
+    assert_eq!(allocation.bytes[0].unwrap(), Char::CapitalB.to_u8());
+    assert_eq!(allocation.bytes[1].unwrap(), Char::SmallA.to_u8());
+    assert_eq!(allocation.bytes[2].unwrap(), Char::SmallR.to_u8());
+}
+
+// Use internal API to find a function in a crate.
+fn get_item<'a>(
+    items: &'a CrateItems,
+    item: (ItemKind, &str),
+) -> Option<&'a stable_mir::CrateItem> {
+    items.iter().find(|crate_item| {
+        (item.0 == crate_item.kind()) && crate_item.name() == item.1
+    })
+}
+
+/// This test will generate and analyze a dummy crate using the stable mir.
+/// For that, it will first write the dummy crate into a file.
+/// Then it will create a `StableMir` using custom arguments and then
+/// it will run the compiler.
+fn main() {
+    let path = "alloc_input.rs";
+    generate_input(&path).unwrap();
+    let args = vec![
+        "rustc".to_string(),
+        "--crate-name".to_string(),
+        CRATE_NAME.to_string(),
+        path.to_string(),
+    ];
+    run!(args, tcx, test_stable_mir(tcx)).unwrap();
+}
+
+fn generate_input(path: &str) -> std::io::Result<()> {
+    let mut file = std::fs::File::create(path)?;
+    write!(
+        file,
+        r#"
+    static FOO: [&str; 2] = ["hi", "there"];
+    static BAR: &str = "Bar";
+
+    pub fn main() {{
+        println!("{{FOO:?}}! {{BAR}}");
+    }}"#
+    )?;
+    Ok(())
+}