rustc_smir/rustc_internal/
mod.rs
1use std::cell::{Cell, RefCell};
7use std::fmt::Debug;
8use std::hash::Hash;
9use std::ops::Index;
10
11use rustc_data_structures::fx;
12use rustc_data_structures::fx::FxIndexMap;
13use rustc_middle::mir::interpret::AllocId;
14use rustc_middle::ty;
15use rustc_middle::ty::TyCtxt;
16use rustc_span::Span;
17use rustc_span::def_id::{CrateNum, DefId};
18use scoped_tls::scoped_thread_local;
19use stable_mir::Error;
20use stable_mir::abi::Layout;
21use stable_mir::ty::IndexedVal;
22
23use crate::rustc_smir::context::TablesWrapper;
24use crate::rustc_smir::{Stable, Tables};
25
26mod internal;
27pub mod pretty;
28
29pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
41 with_tables(|tables| item.stable(tables))
42}
43
44pub fn internal<'tcx, S>(tcx: TyCtxt<'tcx>, item: S) -> S::T<'tcx>
56where
57 S: RustcInternal,
58{
59 with_tables(|tables| item.internal(tables, tcx))
61}
62
63impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
64 type Output = DefId;
65
66 #[inline(always)]
67 fn index(&self, index: stable_mir::DefId) -> &Self::Output {
68 &self.def_ids[index]
69 }
70}
71
72impl<'tcx> Index<stable_mir::ty::Span> for Tables<'tcx> {
73 type Output = Span;
74
75 #[inline(always)]
76 fn index(&self, index: stable_mir::ty::Span) -> &Self::Output {
77 &self.spans[index]
78 }
79}
80
81impl<'tcx> Tables<'tcx> {
82 pub fn crate_item(&mut self, did: DefId) -> stable_mir::CrateItem {
83 stable_mir::CrateItem(self.create_def_id(did))
84 }
85
86 pub fn adt_def(&mut self, did: DefId) -> stable_mir::ty::AdtDef {
87 stable_mir::ty::AdtDef(self.create_def_id(did))
88 }
89
90 pub fn foreign_module_def(&mut self, did: DefId) -> stable_mir::ty::ForeignModuleDef {
91 stable_mir::ty::ForeignModuleDef(self.create_def_id(did))
92 }
93
94 pub fn foreign_def(&mut self, did: DefId) -> stable_mir::ty::ForeignDef {
95 stable_mir::ty::ForeignDef(self.create_def_id(did))
96 }
97
98 pub fn fn_def(&mut self, did: DefId) -> stable_mir::ty::FnDef {
99 stable_mir::ty::FnDef(self.create_def_id(did))
100 }
101
102 pub fn closure_def(&mut self, did: DefId) -> stable_mir::ty::ClosureDef {
103 stable_mir::ty::ClosureDef(self.create_def_id(did))
104 }
105
106 pub fn coroutine_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineDef {
107 stable_mir::ty::CoroutineDef(self.create_def_id(did))
108 }
109
110 pub fn coroutine_closure_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineClosureDef {
111 stable_mir::ty::CoroutineClosureDef(self.create_def_id(did))
112 }
113
114 pub fn alias_def(&mut self, did: DefId) -> stable_mir::ty::AliasDef {
115 stable_mir::ty::AliasDef(self.create_def_id(did))
116 }
117
118 pub fn param_def(&mut self, did: DefId) -> stable_mir::ty::ParamDef {
119 stable_mir::ty::ParamDef(self.create_def_id(did))
120 }
121
122 pub fn br_named_def(&mut self, did: DefId) -> stable_mir::ty::BrNamedDef {
123 stable_mir::ty::BrNamedDef(self.create_def_id(did))
124 }
125
126 pub fn trait_def(&mut self, did: DefId) -> stable_mir::ty::TraitDef {
127 stable_mir::ty::TraitDef(self.create_def_id(did))
128 }
129
130 pub fn generic_def(&mut self, did: DefId) -> stable_mir::ty::GenericDef {
131 stable_mir::ty::GenericDef(self.create_def_id(did))
132 }
133
134 pub fn const_def(&mut self, did: DefId) -> stable_mir::ty::ConstDef {
135 stable_mir::ty::ConstDef(self.create_def_id(did))
136 }
137
138 pub fn impl_def(&mut self, did: DefId) -> stable_mir::ty::ImplDef {
139 stable_mir::ty::ImplDef(self.create_def_id(did))
140 }
141
142 pub fn region_def(&mut self, did: DefId) -> stable_mir::ty::RegionDef {
143 stable_mir::ty::RegionDef(self.create_def_id(did))
144 }
145
146 pub fn coroutine_witness_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineWitnessDef {
147 stable_mir::ty::CoroutineWitnessDef(self.create_def_id(did))
148 }
149
150 pub fn prov(&mut self, aid: AllocId) -> stable_mir::ty::Prov {
151 stable_mir::ty::Prov(self.create_alloc_id(aid))
152 }
153
154 pub(crate) fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
155 self.def_ids.create_or_fetch(did)
156 }
157
158 pub(crate) fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::mir::alloc::AllocId {
159 self.alloc_ids.create_or_fetch(aid)
160 }
161
162 pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
163 self.spans.create_or_fetch(span)
164 }
165
166 pub(crate) fn instance_def(
167 &mut self,
168 instance: ty::Instance<'tcx>,
169 ) -> stable_mir::mir::mono::InstanceDef {
170 self.instances.create_or_fetch(instance)
171 }
172
173 pub(crate) fn static_def(&mut self, did: DefId) -> stable_mir::mir::mono::StaticDef {
174 stable_mir::mir::mono::StaticDef(self.create_def_id(did))
175 }
176
177 pub(crate) fn layout_id(&mut self, layout: rustc_abi::Layout<'tcx>) -> Layout {
178 self.layouts.create_or_fetch(layout)
179 }
180}
181
182pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
183 item.id.into()
184}
185
186scoped_thread_local! (static TLV: Cell<*const ()>);
189
190pub(crate) fn init<'tcx, F, T>(tables: &TablesWrapper<'tcx>, f: F) -> T
191where
192 F: FnOnce() -> T,
193{
194 assert!(!TLV.is_set());
195 let ptr = tables as *const _ as *const ();
196 TLV.set(&Cell::new(ptr), || f())
197}
198
199pub(crate) fn with_tables<R>(f: impl for<'tcx> FnOnce(&mut Tables<'tcx>) -> R) -> R {
202 assert!(TLV.is_set());
203 TLV.with(|tlv| {
204 let ptr = tlv.get();
205 assert!(!ptr.is_null());
206 let wrapper = ptr as *const TablesWrapper<'_>;
207 let mut tables = unsafe { (*wrapper).0.borrow_mut() };
208 f(&mut *tables)
209 })
210}
211
212pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
213where
214 F: FnOnce() -> T,
215{
216 let tables = TablesWrapper(RefCell::new(Tables {
217 tcx,
218 def_ids: IndexMap::default(),
219 alloc_ids: IndexMap::default(),
220 spans: IndexMap::default(),
221 types: IndexMap::default(),
222 instances: IndexMap::default(),
223 ty_consts: IndexMap::default(),
224 mir_consts: IndexMap::default(),
225 layouts: IndexMap::default(),
226 }));
227 stable_mir::compiler_interface::run(&tables, || init(&tables, f))
228}
229
230#[macro_export]
277macro_rules! run {
278 ($args:expr, $callback_fn:ident) => {
279 run_driver!($args, || $callback_fn())
280 };
281 ($args:expr, $callback:expr) => {
282 run_driver!($args, $callback)
283 };
284}
285
286#[macro_export]
291macro_rules! run_with_tcx {
292 ($args:expr, $callback_fn:ident) => {
293 run_driver!($args, |tcx| $callback_fn(tcx), with_tcx)
294 };
295 ($args:expr, $callback:expr) => {
296 run_driver!($args, $callback, with_tcx)
297 };
298}
299
300#[macro_export]
302#[doc(hidden)]
303macro_rules! optional {
304 (with_tcx $ident:ident) => {
305 $ident
306 };
307}
308
309#[macro_export]
316#[doc(hidden)]
317macro_rules! run_driver {
318 ($args:expr, $callback:expr $(, $with_tcx:ident)?) => {{
319 use rustc_driver::{Callbacks, Compilation, run_compiler};
320 use rustc_middle::ty::TyCtxt;
321 use rustc_interface::interface;
322 use stable_mir::CompilerError;
323 use std::ops::ControlFlow;
324
325 pub struct StableMir<B = (), C = (), F = fn($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C>>
326 where
327 B: Send,
328 C: Send,
329 F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
330 {
331 args: Vec<String>,
332 callback: Option<F>,
333 result: Option<ControlFlow<B, C>>,
334 }
335
336 impl<B, C, F> StableMir<B, C, F>
337 where
338 B: Send,
339 C: Send,
340 F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
341 {
342 pub fn new(args: Vec<String>, callback: F) -> Self {
344 StableMir { args, callback: Some(callback), result: None }
345 }
346
347 pub fn run(&mut self) -> Result<C, CompilerError<B>> {
349 let compiler_result = rustc_driver::catch_fatal_errors(|| -> interface::Result::<()> {
350 run_compiler(&self.args.clone(), self);
351 Ok(())
352 });
353 match (compiler_result, self.result.take()) {
354 (Ok(Ok(())), Some(ControlFlow::Continue(value))) => Ok(value),
355 (Ok(Ok(())), Some(ControlFlow::Break(value))) => {
356 Err(CompilerError::Interrupted(value))
357 }
358 (Ok(Ok(_)), None) => Err(CompilerError::Skipped),
359 (Ok(Err(_)), _) | (Err(_), _) => Err(CompilerError::Failed),
367 }
368 }
369 }
370
371 impl<B, C, F> Callbacks for StableMir<B, C, F>
372 where
373 B: Send,
374 C: Send,
375 F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
376 {
377 fn after_analysis<'tcx>(
380 &mut self,
381 _compiler: &interface::Compiler,
382 tcx: TyCtxt<'tcx>,
383 ) -> Compilation {
384 if let Some(callback) = self.callback.take() {
385 rustc_internal::run(tcx, || {
386 self.result = Some(callback($(optional!($with_tcx tcx))?));
387 })
388 .unwrap();
389 if self.result.as_ref().is_some_and(|val| val.is_continue()) {
390 Compilation::Continue
391 } else {
392 Compilation::Stop
393 }
394 } else {
395 Compilation::Continue
396 }
397 }
398 }
399
400 StableMir::new($args, $callback).run()
401 }};
402}
403
404pub struct IndexMap<K, V> {
407 index_map: fx::FxIndexMap<K, V>,
408}
409
410impl<K, V> Default for IndexMap<K, V> {
411 fn default() -> Self {
412 Self { index_map: FxIndexMap::default() }
413 }
414}
415
416impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> IndexMap<K, V> {
417 pub fn create_or_fetch(&mut self, key: K) -> V {
418 let len = self.index_map.len();
419 let v = self.index_map.entry(key).or_insert(V::to_val(len));
420 *v
421 }
422}
423
424impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> Index<V>
425 for IndexMap<K, V>
426{
427 type Output = K;
428
429 fn index(&self, index: V) -> &Self::Output {
430 let (k, v) = self.index_map.get_index(index.to_index()).unwrap();
431 assert_eq!(*v, index, "Provided value doesn't match with indexed value");
432 k
433 }
434}
435
436pub trait RustcInternal {
440 type T<'tcx>;
441 fn internal<'tcx>(&self, tables: &mut Tables<'_>, tcx: TyCtxt<'tcx>) -> Self::T<'tcx>;
442}