Add support to global allocation to stable-mir

Author: celinval

compiler/rustc_smir/src/rustc_internal/internal.rs

@@ -7,12 +7,13 @@
 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::{
     Binder, BoundRegionKind, BoundTyKind, BoundVariableKind, ClosureKind, Const, GenericArgKind,
     GenericArgs, Region, TraitRef, Ty,
 };
-use stable_mir::{AllocId, CrateItem, DefId};
+use stable_mir::{CrateItem, DefId};
 
 use super::RustcInternal;
 

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 {
+    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,14 +17,11 @@ use rustc_middle::mir::mono::MonoItem;
 use rustc_middle::ty::{self, Instance, ParamEnv, 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, CopyNonOverlapping, Statement, UserTypeProjection, VariantIdx};
-use stable_mir::ty::{
-    AdtDef, AdtKind, ClosureDef, ClosureKind, Const, ConstId, ConstantKind, EarlyParamRegion,
-    FloatTy, FnDef, GenericArgs, GenericParamDef, IntTy, LineInfo, Movability, RigidTy, Span,
-    TyKind, UintTy,
-};
-use stable_mir::{self, opaque, Context, CrateItem, Filename, ItemKind};
+use stable_mir::ty::{AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind, Const, ConstId, ConstantKind, FloatTy, FnDef, GenericArgs, GenericParamDef, IntTy, LineInfo, Movability, RigidTy, Span, TyKind, UintTy, EarlyParamRegion};
+use stable_mir::{self, opaque, Context, CrateItem, Error, Filename, ItemKind};
 use std::cell::RefCell;
 use tracing::debug;
 
@@ -91,6 +88,11 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         new_item_kind(tables.tcx.def_kind(tables[item.0]))
     }
 
+    fn is_foreign_item(&self, item: CrateItem) -> bool {
+        let tables = self.0.borrow();
+        tables.tcx.is_foreign_item(tables[item.0])
+    }
+
     fn adt_kind(&self, def: AdtDef) -> AdtKind {
         let mut tables = self.0.borrow_mut();
         let ty = tables.tcx.type_of(def.0.internal(&mut *tables)).instantiate_identity().kind();
@@ -302,14 +304,26 @@ impl<'tcx> Context for TablesWrapper<'tcx> {
         let closure_kind = kind.internal(&mut *tables);
         Instance::resolve_closure(tables.tcx, def_id, args_ref, closure_kind).stable(&mut *tables)
     }
+
+    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)
+    }
 }
 
 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>,
@@ -1507,6 +1521,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;
 
@@ -1519,6 +1541,25 @@ impl<'tcx> Stable<'tcx> for mir::interpret::Allocation {
     }
 }
 
+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 {
@@ -1906,6 +1947,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>,
@@ -1927,3 +1976,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)),
+        }
+    }
+}

compiler/stable_mir/src/error.rs

@@ -27,7 +27,7 @@ pub enum CompilerError<T> {
 pub struct Error(String);
 
 impl Error {
-    pub(crate) fn new(msg: String) -> Self {
+    pub fn new(msg: String) -> Self {
         Self(msg)
     }
 }

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;
@@ -36,9 +36,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};
+use crate::ty::{AdtDef, AdtKind, Allocation, ClosureDef, ClosureKind};
 pub use error::*;
 use mir::mono::Instance;
 use ty::{Const, FnDef, GenericArgs};
@@ -72,19 +73,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>;
 
@@ -140,6 +128,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)
@@ -189,6 +181,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.
@@ -224,6 +218,9 @@ pub trait Context {
     /// Returns the `kind` of given `DefId`
     fn item_kind(&self, item: CrateItem) -> ItemKind;
 
+    /// Returns whether this is a foreign item.
+    fn is_foreign_item(&self, item: CrateItem) -> bool;
+
     /// Returns the kind of a given algebraic data type
     fn adt_kind(&self, def: AdtDef) -> AdtKind;
 
@@ -275,6 +272,12 @@ 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;
 }
 
 // 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,37 @@
+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.
+    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))
+    }
+}
+
+/// 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;
 
@@ -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))
@@ -119,6 +124,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);
 
@@ -138,9 +155,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 {

compiler/stable_mir/src/ty.rs

@@ -1,8 +1,9 @@
 use super::{
     mir::Safety,
     mir::{Body, Mutability},
-    with, AllocId, DefId, Symbol,
+    with, DefId, Symbol,
 };
+use crate::mir::alloc::AllocId;
 use crate::{Filename, Opaque};
 use std::fmt::{self, Debug, Display, Formatter};