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.

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 @@ fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
     }
 }
 
+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 @@ fn internal(&self, tables: &mut Tables<'tcx>) -> Self::T {
         (*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))
+    }
+}

compiler/rustc_smir/src/rustc_internal/mod.rs

@@ -118,7 +118,7 @@ pub(crate) fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
         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)
     }
 

compiler/rustc_smir/src/rustc_smir/mod.rs

@@ -17,15 +17,16 @@
 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 @@ fn eval_target_usize(&self, cnst: &Const) -> Result<u64, Error> {
             .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 @@ fn new_rigid_ty(&self, kind: RigidTy) -> stable_mir::ty::Ty {
 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 @@ fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
     }
 }
 
+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 @@ fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
     }
 }
 
+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 @@ fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
     }
 }
 
+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 @@ fn stable(&self, tables: &mut Tables<'tcx>) -> Self::T {
         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)),
+        }
+    }
+}

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 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 @@ fn to_index(&self) -> usize {
     }
 }
 
-/// 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 @@ pub fn ty(&self) -> Ty {
         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 @@ fn resolve_closure(
         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

compiler/stable_mir/src/mir.rs

@@ -1,3 +1,4 @@
+pub mod alloc;
 mod body;
 pub mod mono;
 pub mod pretty;

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
+    }
+}

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 @@ pub fn body(&self) -> Option<Body> {
         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 @@ fn from(value: Instance) -> Self {
     }
 }
 
+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 @@ fn try_from(value: CrateItem) -> Result<Self, Self::Error> {
 }
 
 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 {

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};
 

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(())
+}