rustc_smir/rustc_internal/
mod.rs1use 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::compiler_interface::SmirInterface;
22use stable_mir::ty::IndexedVal;
23
24use crate::rustc_smir::context::SmirCtxt;
25use crate::rustc_smir::{Stable, Tables};
26use crate::stable_mir;
27
28mod internal;
29pub mod pretty;
30
31pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
43 with_tables(|tables| item.stable(tables))
44}
45
46pub fn internal<'tcx, S>(tcx: TyCtxt<'tcx>, item: S) -> S::T<'tcx>
58where
59 S: RustcInternal,
60{
61 with_tables(|tables| item.internal(tables, tcx))
63}
64
65impl<'tcx> Index<stable_mir::DefId> for Tables<'tcx> {
66 type Output = DefId;
67
68 #[inline(always)]
69 fn index(&self, index: stable_mir::DefId) -> &Self::Output {
70 &self.def_ids[index]
71 }
72}
73
74impl<'tcx> Index<stable_mir::ty::Span> for Tables<'tcx> {
75 type Output = Span;
76
77 #[inline(always)]
78 fn index(&self, index: stable_mir::ty::Span) -> &Self::Output {
79 &self.spans[index]
80 }
81}
82
83impl<'tcx> Tables<'tcx> {
84 pub fn crate_item(&mut self, did: DefId) -> stable_mir::CrateItem {
85 stable_mir::CrateItem(self.create_def_id(did))
86 }
87
88 pub fn adt_def(&mut self, did: DefId) -> stable_mir::ty::AdtDef {
89 stable_mir::ty::AdtDef(self.create_def_id(did))
90 }
91
92 pub fn foreign_module_def(&mut self, did: DefId) -> stable_mir::ty::ForeignModuleDef {
93 stable_mir::ty::ForeignModuleDef(self.create_def_id(did))
94 }
95
96 pub fn foreign_def(&mut self, did: DefId) -> stable_mir::ty::ForeignDef {
97 stable_mir::ty::ForeignDef(self.create_def_id(did))
98 }
99
100 pub fn fn_def(&mut self, did: DefId) -> stable_mir::ty::FnDef {
101 stable_mir::ty::FnDef(self.create_def_id(did))
102 }
103
104 pub fn closure_def(&mut self, did: DefId) -> stable_mir::ty::ClosureDef {
105 stable_mir::ty::ClosureDef(self.create_def_id(did))
106 }
107
108 pub fn coroutine_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineDef {
109 stable_mir::ty::CoroutineDef(self.create_def_id(did))
110 }
111
112 pub fn coroutine_closure_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineClosureDef {
113 stable_mir::ty::CoroutineClosureDef(self.create_def_id(did))
114 }
115
116 pub fn alias_def(&mut self, did: DefId) -> stable_mir::ty::AliasDef {
117 stable_mir::ty::AliasDef(self.create_def_id(did))
118 }
119
120 pub fn param_def(&mut self, did: DefId) -> stable_mir::ty::ParamDef {
121 stable_mir::ty::ParamDef(self.create_def_id(did))
122 }
123
124 pub fn br_named_def(&mut self, did: DefId) -> stable_mir::ty::BrNamedDef {
125 stable_mir::ty::BrNamedDef(self.create_def_id(did))
126 }
127
128 pub fn trait_def(&mut self, did: DefId) -> stable_mir::ty::TraitDef {
129 stable_mir::ty::TraitDef(self.create_def_id(did))
130 }
131
132 pub fn generic_def(&mut self, did: DefId) -> stable_mir::ty::GenericDef {
133 stable_mir::ty::GenericDef(self.create_def_id(did))
134 }
135
136 pub fn const_def(&mut self, did: DefId) -> stable_mir::ty::ConstDef {
137 stable_mir::ty::ConstDef(self.create_def_id(did))
138 }
139
140 pub fn impl_def(&mut self, did: DefId) -> stable_mir::ty::ImplDef {
141 stable_mir::ty::ImplDef(self.create_def_id(did))
142 }
143
144 pub fn region_def(&mut self, did: DefId) -> stable_mir::ty::RegionDef {
145 stable_mir::ty::RegionDef(self.create_def_id(did))
146 }
147
148 pub fn coroutine_witness_def(&mut self, did: DefId) -> stable_mir::ty::CoroutineWitnessDef {
149 stable_mir::ty::CoroutineWitnessDef(self.create_def_id(did))
150 }
151
152 pub fn assoc_def(&mut self, did: DefId) -> stable_mir::ty::AssocDef {
153 stable_mir::ty::AssocDef(self.create_def_id(did))
154 }
155
156 pub fn opaque_def(&mut self, did: DefId) -> stable_mir::ty::OpaqueDef {
157 stable_mir::ty::OpaqueDef(self.create_def_id(did))
158 }
159
160 pub fn prov(&mut self, aid: AllocId) -> stable_mir::ty::Prov {
161 stable_mir::ty::Prov(self.create_alloc_id(aid))
162 }
163
164 pub(crate) fn create_def_id(&mut self, did: DefId) -> stable_mir::DefId {
165 self.def_ids.create_or_fetch(did)
166 }
167
168 pub(crate) fn create_alloc_id(&mut self, aid: AllocId) -> stable_mir::mir::alloc::AllocId {
169 self.alloc_ids.create_or_fetch(aid)
170 }
171
172 pub(crate) fn create_span(&mut self, span: Span) -> stable_mir::ty::Span {
173 self.spans.create_or_fetch(span)
174 }
175
176 pub(crate) fn instance_def(
177 &mut self,
178 instance: ty::Instance<'tcx>,
179 ) -> stable_mir::mir::mono::InstanceDef {
180 self.instances.create_or_fetch(instance)
181 }
182
183 pub(crate) fn static_def(&mut self, did: DefId) -> stable_mir::mir::mono::StaticDef {
184 stable_mir::mir::mono::StaticDef(self.create_def_id(did))
185 }
186
187 pub(crate) fn layout_id(&mut self, layout: rustc_abi::Layout<'tcx>) -> Layout {
188 self.layouts.create_or_fetch(layout)
189 }
190}
191
192pub fn crate_num(item: &stable_mir::Crate) -> CrateNum {
193 item.id.into()
194}
195
196scoped_thread_local! (static TLV: Cell<*const ()>);
199
200pub(crate) fn init<'tcx, F, T>(cx: &SmirCtxt<'tcx>, f: F) -> T
201where
202 F: FnOnce() -> T,
203{
204 assert!(!TLV.is_set());
205 let ptr = cx as *const _ as *const ();
206 TLV.set(&Cell::new(ptr), || f())
207}
208
209pub(crate) fn with_tables<R>(f: impl for<'tcx> FnOnce(&mut Tables<'tcx>) -> R) -> R {
212 assert!(TLV.is_set());
213 TLV.with(|tlv| {
214 let ptr = tlv.get();
215 assert!(!ptr.is_null());
216 let context = ptr as *const SmirCtxt<'_>;
217 let mut tables = unsafe { (*context).0.borrow_mut() };
218 f(&mut *tables)
219 })
220}
221
222pub fn run<F, T>(tcx: TyCtxt<'_>, f: F) -> Result<T, Error>
223where
224 F: FnOnce() -> T,
225{
226 let tables = SmirCtxt(RefCell::new(Tables {
227 tcx,
228 def_ids: IndexMap::default(),
229 alloc_ids: IndexMap::default(),
230 spans: IndexMap::default(),
231 types: IndexMap::default(),
232 instances: IndexMap::default(),
233 ty_consts: IndexMap::default(),
234 mir_consts: IndexMap::default(),
235 layouts: IndexMap::default(),
236 }));
237
238 let interface = SmirInterface { cx: tables };
239
240 stable_mir::compiler_interface::run(&interface, || init(&interface.cx, f))
243}
244
245#[macro_export]
294macro_rules! run {
295 ($args:expr, $callback_fn:ident) => {
296 run_driver!($args, || $callback_fn())
297 };
298 ($args:expr, $callback:expr) => {
299 run_driver!($args, $callback)
300 };
301}
302
303#[macro_export]
308macro_rules! run_with_tcx {
309 ($args:expr, $callback_fn:ident) => {
310 run_driver!($args, |tcx| $callback_fn(tcx), with_tcx)
311 };
312 ($args:expr, $callback:expr) => {
313 run_driver!($args, $callback, with_tcx)
314 };
315}
316
317#[macro_export]
319#[doc(hidden)]
320macro_rules! optional {
321 (with_tcx $ident:ident) => {
322 $ident
323 };
324}
325
326#[macro_export]
333#[doc(hidden)]
334macro_rules! run_driver {
335 ($args:expr, $callback:expr $(, $with_tcx:ident)?) => {{
336 use rustc_driver::{Callbacks, Compilation, run_compiler};
337 use rustc_middle::ty::TyCtxt;
338 use rustc_interface::interface;
339 use rustc_smir::rustc_internal;
340 use stable_mir::CompilerError;
341 use std::ops::ControlFlow;
342
343 pub struct StableMir<B = (), C = (), F = fn($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C>>
344 where
345 B: Send,
346 C: Send,
347 F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
348 {
349 callback: Option<F>,
350 result: Option<ControlFlow<B, C>>,
351 }
352
353 impl<B, C, F> StableMir<B, C, F>
354 where
355 B: Send,
356 C: Send,
357 F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
358 {
359 pub fn new(callback: F) -> Self {
361 StableMir { callback: Some(callback), result: None }
362 }
363
364 pub fn run(&mut self, args: &[String]) -> Result<C, CompilerError<B>> {
366 let compiler_result = rustc_driver::catch_fatal_errors(|| -> interface::Result::<()> {
367 run_compiler(&args, self);
368 Ok(())
369 });
370 match (compiler_result, self.result.take()) {
371 (Ok(Ok(())), Some(ControlFlow::Continue(value))) => Ok(value),
372 (Ok(Ok(())), Some(ControlFlow::Break(value))) => {
373 Err(CompilerError::Interrupted(value))
374 }
375 (Ok(Ok(_)), None) => Err(CompilerError::Skipped),
376 (Ok(Err(_)), _) | (Err(_), _) => Err(CompilerError::Failed),
384 }
385 }
386 }
387
388 impl<B, C, F> Callbacks for StableMir<B, C, F>
389 where
390 B: Send,
391 C: Send,
392 F: FnOnce($(optional!($with_tcx TyCtxt))?) -> ControlFlow<B, C> + Send,
393 {
394 fn after_analysis<'tcx>(
397 &mut self,
398 _compiler: &interface::Compiler,
399 tcx: TyCtxt<'tcx>,
400 ) -> Compilation {
401 if let Some(callback) = self.callback.take() {
402 rustc_internal::run(tcx, || {
403 self.result = Some(callback($(optional!($with_tcx tcx))?));
404 })
405 .unwrap();
406 if self.result.as_ref().is_some_and(|val| val.is_continue()) {
407 Compilation::Continue
408 } else {
409 Compilation::Stop
410 }
411 } else {
412 Compilation::Continue
413 }
414 }
415 }
416
417 StableMir::new($callback).run($args)
418 }};
419}
420
421pub struct IndexMap<K, V> {
424 index_map: fx::FxIndexMap<K, V>,
425}
426
427impl<K, V> Default for IndexMap<K, V> {
428 fn default() -> Self {
429 Self { index_map: FxIndexMap::default() }
430 }
431}
432
433impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> IndexMap<K, V> {
434 pub fn create_or_fetch(&mut self, key: K) -> V {
435 let len = self.index_map.len();
436 let v = self.index_map.entry(key).or_insert(V::to_val(len));
437 *v
438 }
439}
440
441impl<K: PartialEq + Hash + Eq, V: Copy + Debug + PartialEq + IndexedVal> Index<V>
442 for IndexMap<K, V>
443{
444 type Output = K;
445
446 fn index(&self, index: V) -> &Self::Output {
447 let (k, v) = self.index_map.get_index(index.to_index()).unwrap();
448 assert_eq!(*v, index, "Provided value doesn't match with indexed value");
449 k
450 }
451}
452
453pub trait RustcInternal {
457 type T<'tcx>;
458 fn internal<'tcx>(&self, tables: &mut Tables<'_>, tcx: TyCtxt<'tcx>) -> Self::T<'tcx>;
459}