Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for global allocation in smir #118012

Merged
merged 3 commits into from
Nov 22, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 26 additions & 3 deletions compiler/rustc_smir/src/rustc_internal/internal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -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>;

Expand Down Expand Up @@ -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))
}
}
2 changes: 1 addition & 1 deletion compiler/rustc_smir/src/rustc_internal/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}

Expand Down
92 changes: 87 additions & 5 deletions compiler/rustc_smir/src/rustc_smir/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand All @@ -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>,
Expand Down Expand Up @@ -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;

Expand All @@ -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 {
Expand Down Expand Up @@ -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>,
Expand All @@ -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)),
}
}
}
33 changes: 18 additions & 15 deletions compiler/stable_mir/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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};
Expand Down Expand Up @@ -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>;

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions compiler/stable_mir/src/mir.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
pub mod alloc;
mod body;
pub mod mono;
pub mod pretty;
Expand Down
51 changes: 51 additions & 0 deletions compiler/stable_mir/src/mir/alloc.rs
Original file line number Diff line number Diff line change
@@ -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
}
}
25 changes: 24 additions & 1 deletion compiler/stable_mir/src/mir/mono.rs
Original file line number Diff line number Diff line change
@@ -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};

Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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);

Expand All @@ -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 {
Expand Down
Loading