Skip to content

Commit b43e3b9

Browse files
oli-obkoli-obk
committed
Add type folder to SMIR
1 parent 3767e31 commit b43e3b9

File tree

5 files changed

+255
-6
lines changed

5 files changed

+255
-6
lines changed

‎compiler/rustc_smir/src/rustc_smir/mod.rs

Lines changed: 37 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -103,8 +103,13 @@ impl<'tcx> Context for Tables<'tcx> {
103103
}
104104

105105
fn ty_kind(&mut self, ty: crate::stable_mir::ty::Ty) -> TyKind {
106-
let ty = self.types[ty.0];
107-
ty.stable(self)
106+
self.types[ty.0].clone().stable(self)
107+
}
108+
109+
fn mk_ty(&mut self, kind: TyKind) -> stable_mir::ty::Ty {
110+
let n = self.types.len();
111+
self.types.push(MaybeStable::Stable(kind));
112+
stable_mir::ty::Ty(n)
108113
}
109114

110115
fn generics_of(&mut self, def_id: stable_mir::DefId) -> stable_mir::ty::Generics {
@@ -128,20 +133,47 @@ impl<'tcx> Context for Tables<'tcx> {
128133
}
129134
}
130135

136+
#[derive(Clone)]
137+
pub enum MaybeStable<S, R> {
138+
Stable(S),
139+
Rustc(R),
140+
}
141+
142+
impl<'tcx, S, R> MaybeStable<S, R> {
143+
fn stable(self, tables: &mut Tables<'tcx>) -> S
144+
where
145+
R: Stable<'tcx, T = S>,
146+
{
147+
match self {
148+
MaybeStable::Stable(s) => s,
149+
MaybeStable::Rustc(r) => r.stable(tables),
150+
}
151+
}
152+
}
153+
154+
impl<S, R: PartialEq> PartialEq<R> for MaybeStable<S, R> {
155+
fn eq(&self, other: &R) -> bool {
156+
match self {
157+
MaybeStable::Stable(_) => false,
158+
MaybeStable::Rustc(r) => r == other,
159+
}
160+
}
161+
}
162+
131163
pub struct Tables<'tcx> {
132164
pub tcx: TyCtxt<'tcx>,
133165
pub def_ids: Vec<DefId>,
134166
pub alloc_ids: Vec<AllocId>,
135-
pub types: Vec<Ty<'tcx>>,
167+
pub types: Vec<MaybeStable<stable_mir::ty::TyKind, Ty<'tcx>>>,
136168
}
137169

138170
impl<'tcx> Tables<'tcx> {
139171
fn intern_ty(&mut self, ty: Ty<'tcx>) -> stable_mir::ty::Ty {
140-
if let Some(id) = self.types.iter().position(|&t| t == ty) {
172+
if let Some(id) = self.types.iter().position(|t| *t == ty) {
141173
return stable_mir::ty::Ty(id);
142174
}
143175
let id = self.types.len();
144-
self.types.push(ty);
176+
self.types.push(MaybeStable::Rustc(ty));
145177
stable_mir::ty::Ty(id)
146178
}
147179
}

‎compiler/rustc_smir/src/stable_mir/fold.rs

Lines changed: 207 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,207 @@
1+
use std::ops::ControlFlow;
2+
3+
use crate::rustc_internal::Opaque;
4+
5+
use super::ty::{
6+
Allocation, Binder, Const, ConstDef, ExistentialPredicate, FnSig, GenericArgKind, GenericArgs,
7+
Promoted, RigidTy, TermKind, Ty, UnevaluatedConst,
8+
};
9+
10+
pub trait Folder: Sized {
11+
type Break;
12+
fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::Break, Ty> {
13+
ty.super_fold(self)
14+
}
15+
fn fold_const(&mut self, c: &Const) -> ControlFlow<Self::Break, Const> {
16+
c.super_fold(self)
17+
}
18+
}
19+
20+
pub trait Foldable: Sized + Clone {
21+
fn fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
22+
self.super_fold(folder)
23+
}
24+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self>;
25+
}
26+
27+
impl Foldable for Ty {
28+
fn fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
29+
folder.visit_ty(self)
30+
}
31+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
32+
let mut kind = self.kind();
33+
match &mut kind {
34+
super::ty::TyKind::RigidTy(ty) => *ty = ty.fold(folder)?,
35+
super::ty::TyKind::Alias(_, alias) => alias.args = alias.args.fold(folder)?,
36+
super::ty::TyKind::Param(_) => {}
37+
super::ty::TyKind::Bound(_, _) => {}
38+
}
39+
ControlFlow::Continue(kind.into())
40+
}
41+
}
42+
43+
impl Foldable for Const {
44+
fn fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
45+
folder.fold_const(self)
46+
}
47+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
48+
let mut this = self.clone();
49+
match &mut this.literal {
50+
super::ty::ConstantKind::Allocated(alloc) => *alloc = alloc.fold(folder)?,
51+
super::ty::ConstantKind::Unevaluated(uv) => *uv = uv.fold(folder)?,
52+
super::ty::ConstantKind::ParamCt(param) => *param = param.fold(folder)?,
53+
}
54+
ControlFlow::Continue(this)
55+
}
56+
}
57+
58+
impl Foldable for Opaque {
59+
fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
60+
ControlFlow::Continue(self.clone())
61+
}
62+
}
63+
64+
impl Foldable for Allocation {
65+
fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
66+
ControlFlow::Continue(self.clone())
67+
}
68+
}
69+
70+
impl Foldable for UnevaluatedConst {
71+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
72+
let UnevaluatedConst { ty, def, args, promoted } = self;
73+
ControlFlow::Continue(UnevaluatedConst {
74+
ty: ty.fold(folder)?,
75+
def: def.fold(folder)?,
76+
args: args.fold(folder)?,
77+
promoted: promoted.fold(folder)?,
78+
})
79+
}
80+
}
81+
82+
impl Foldable for ConstDef {
83+
fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
84+
ControlFlow::Continue(self.clone())
85+
}
86+
}
87+
88+
impl<T: Foldable> Foldable for Option<T> {
89+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
90+
ControlFlow::Continue(match self {
91+
Some(val) => Some(val.fold(folder)?),
92+
None => None,
93+
})
94+
}
95+
}
96+
97+
impl Foldable for Promoted {
98+
fn super_fold<V: Folder>(&self, _folder: &mut V) -> ControlFlow<V::Break, Self> {
99+
ControlFlow::Continue(self.clone())
100+
}
101+
}
102+
103+
impl Foldable for GenericArgs {
104+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
105+
ControlFlow::Continue(GenericArgs(self.0.fold(folder)?))
106+
}
107+
}
108+
109+
impl Foldable for GenericArgKind {
110+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
111+
let mut this = self.clone();
112+
match &mut this {
113+
GenericArgKind::Lifetime(lt) => *lt = lt.fold(folder)?,
114+
GenericArgKind::Type(t) => *t = t.fold(folder)?,
115+
GenericArgKind::Const(c) => *c = c.fold(folder)?,
116+
}
117+
ControlFlow::Continue(this)
118+
}
119+
}
120+
121+
impl Foldable for RigidTy {
122+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
123+
let mut this = self.clone();
124+
match &mut this {
125+
RigidTy::Bool
126+
| RigidTy::Char
127+
| RigidTy::Int(_)
128+
| RigidTy::Uint(_)
129+
| RigidTy::Float(_)
130+
| RigidTy::Never
131+
| RigidTy::Foreign(_)
132+
| RigidTy::Str => {}
133+
RigidTy::Array(t, c) => {
134+
*t = t.fold(folder)?;
135+
*c = c.fold(folder)?;
136+
}
137+
RigidTy::Slice(inner) => *inner = inner.fold(folder)?,
138+
RigidTy::RawPtr(ty, _) => *ty = ty.fold(folder)?,
139+
RigidTy::Ref(_, ty, _) => *ty = ty.fold(folder)?,
140+
RigidTy::FnDef(_, args) => *args = args.fold(folder)?,
141+
RigidTy::FnPtr(sig) => *sig = sig.fold(folder)?,
142+
RigidTy::Closure(_, args) => *args = args.fold(folder)?,
143+
RigidTy::Generator(_, args, _) => *args = args.fold(folder)?,
144+
RigidTy::Dynamic(pred, r, _) => {
145+
*pred = pred.fold(folder)?;
146+
*r = r.fold(folder)?;
147+
}
148+
RigidTy::Tuple(fields) => *fields = fields.fold(folder)?,
149+
RigidTy::Adt(_, args) => *args = args.fold(folder)?,
150+
}
151+
ControlFlow::Continue(this)
152+
}
153+
}
154+
155+
impl<T: Foldable> Foldable for Vec<T> {
156+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
157+
let mut this = self.clone();
158+
for arg in &mut this {
159+
*arg = arg.fold(folder)?;
160+
}
161+
ControlFlow::Continue(this)
162+
}
163+
}
164+
165+
impl<T: Foldable> Foldable for Binder<T> {
166+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
167+
ControlFlow::Continue(Self {
168+
value: self.value.fold(folder)?,
169+
bound_vars: self.bound_vars.clone(),
170+
})
171+
}
172+
}
173+
174+
impl Foldable for ExistentialPredicate {
175+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
176+
let mut this = self.clone();
177+
match &mut this {
178+
ExistentialPredicate::Trait(tr) => tr.generic_args = tr.generic_args.fold(folder)?,
179+
ExistentialPredicate::Projection(p) => {
180+
p.term = p.term.fold(folder)?;
181+
p.generic_args = p.generic_args.fold(folder)?;
182+
}
183+
ExistentialPredicate::AutoTrait(_) => {}
184+
}
185+
ControlFlow::Continue(this)
186+
}
187+
}
188+
189+
impl Foldable for TermKind {
190+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
191+
ControlFlow::Continue(match self {
192+
TermKind::Type(t) => TermKind::Type(t.fold(folder)?),
193+
TermKind::Const(c) => TermKind::Const(c.fold(folder)?),
194+
})
195+
}
196+
}
197+
198+
impl Foldable for FnSig {
199+
fn super_fold<V: Folder>(&self, folder: &mut V) -> ControlFlow<V::Break, Self> {
200+
ControlFlow::Continue(Self {
201+
inputs_and_output: self.inputs_and_output.fold(folder)?,
202+
c_variadic: self.c_variadic,
203+
unsafety: self.unsafety,
204+
abi: self.abi.clone(),
205+
})
206+
}
207+
}

‎compiler/rustc_smir/src/stable_mir/mir/body.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -391,7 +391,7 @@ pub enum Mutability {
391391
Mut,
392392
}
393393

394-
#[derive(Clone, Debug)]
394+
#[derive(Copy, Clone, Debug)]
395395
pub enum Safety {
396396
Unsafe,
397397
Normal,

‎compiler/rustc_smir/src/stable_mir/mod.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ use self::ty::{
2020
};
2121
use crate::rustc_smir::Tables;
2222

23+
pub mod fold;
2324
pub mod mir;
2425
pub mod ty;
2526
pub mod visitor;
@@ -158,6 +159,9 @@ pub trait Context {
158159
/// Obtain the representation of a type.
159160
fn ty_kind(&mut self, ty: Ty) -> TyKind;
160161

162+
/// Create a new `Ty` from scratch without information from rustc.
163+
fn mk_ty(&mut self, kind: TyKind) -> Ty;
164+
161165
/// HACK: Until we have fully stable consumers, we need an escape hatch
162166
/// to get `DefId`s out of `CrateItem`s.
163167
fn rustc_tables(&mut self, f: &mut dyn FnMut(&mut Tables<'_>));

‎compiler/rustc_smir/src/stable_mir/ty.rs

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,12 @@ impl Ty {
1010
}
1111
}
1212

13+
impl From<TyKind> for Ty {
14+
fn from(value: TyKind) -> Self {
15+
with(|context| context.mk_ty(value))
16+
}
17+
}
18+
1319
#[derive(Debug, Clone)]
1420
pub struct Const {
1521
pub literal: ConstantKind,

0 commit comments

Comments
 (0)