rustc_smir/rustc_internal/
mod.rs

1//! Module that implements the bridge between Stable MIR and internal compiler MIR.
2//!
3//! For that, we define APIs that will temporarily be public to 3P that exposes rustc internal APIs
4//! until stable MIR is complete.
5
6use 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
29/// Convert an internal Rust compiler item into its stable counterpart, if one exists.
30///
31/// # Warning
32///
33/// This function is unstable, and its behavior may change at any point.
34/// E.g.: Items that were previously supported, may no longer be supported, or its translation may
35/// change.
36///
37/// # Panics
38///
39/// This function will panic if StableMIR has not been properly initialized.
40pub fn stable<'tcx, S: Stable<'tcx>>(item: S) -> S::T {
41    with_tables(|tables| item.stable(tables))
42}
43
44/// Convert a stable item into its internal Rust compiler counterpart, if one exists.
45///
46/// # Warning
47///
48/// This function is unstable, and it's behavior may change at any point.
49/// Not every stable item can be converted to an internal one.
50/// Furthermore, items that were previously supported, may no longer be supported in newer versions.
51///
52/// # Panics
53///
54/// This function will panic if StableMIR has not been properly initialized.
55pub fn internal<'tcx, S>(tcx: TyCtxt<'tcx>, item: S) -> S::T<'tcx>
56where
57    S: RustcInternal,
58{
59    // The tcx argument ensures that the item won't outlive the type context.
60    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
186// A thread local variable that stores a pointer to the tables mapping between TyCtxt
187// datastructures and stable MIR datastructures
188scoped_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
199/// Loads the current context and calls a function with it.
200/// Do not nest these, as that will ICE.
201pub(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/// Instantiate and run the compiler with the provided arguments and callback.
231///
232/// The callback will be invoked after the compiler ran all its analyses, but before code generation.
233/// Note that this macro accepts two different formats for the callback:
234/// 1. An ident that resolves to a function that accepts no argument and returns `ControlFlow<B, C>`
235/// ```ignore(needs-extern-crate)
236/// # extern crate rustc_driver;
237/// # extern crate rustc_interface;
238/// # #[macro_use]
239/// # extern crate rustc_smir;
240/// # extern crate stable_mir;
241/// #
242/// # fn main() {
243/// #   use std::ops::ControlFlow;
244/// #   use stable_mir::CompilerError;
245///     fn analyze_code() -> ControlFlow<(), ()> {
246///         // Your code goes in here.
247/// #       ControlFlow::Continue(())
248///     }
249/// #   let args = vec!["--verbose".to_string()];
250///     let result = run!(args, analyze_code);
251/// #   assert_eq!(result, Err(CompilerError::Skipped))
252/// # }
253/// ```
254/// 2. A closure expression:
255/// ```ignore(needs-extern-crate)
256/// # extern crate rustc_driver;
257/// # extern crate rustc_interface;
258/// # #[macro_use]
259/// # extern crate rustc_smir;
260/// # extern crate stable_mir;
261/// #
262/// # fn main() {
263/// #   use std::ops::ControlFlow;
264/// #   use stable_mir::CompilerError;
265///     fn analyze_code(extra_args: Vec<String>) -> ControlFlow<(), ()> {
266/// #       let _ = extra_args;
267///         // Your code goes in here.
268/// #       ControlFlow::Continue(())
269///     }
270/// #   let args = vec!["--verbose".to_string()];
271/// #   let extra_args = vec![];
272///     let result = run!(args, || analyze_code(extra_args));
273/// #   assert_eq!(result, Err(CompilerError::Skipped))
274/// # }
275/// ```
276#[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/// Instantiate and run the compiler with the provided arguments and callback.
287///
288/// This is similar to `run` but it invokes the callback with the compiler's `TyCtxt`,
289/// which can be used to invoke internal APIs.
290#[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/// Optionally include an ident. This is needed due to macro hygiene.
301#[macro_export]
302#[doc(hidden)]
303macro_rules! optional {
304    (with_tcx $ident:ident) => {
305        $ident
306    };
307}
308
309/// Prefer using [run!] and [run_with_tcx] instead.
310///
311/// This macro implements the instantiation of a StableMIR driver, and it will invoke
312/// the given callback after the compiler analyses.
313///
314/// The third argument determines whether the callback requires `tcx` as an argument.
315#[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            /// Creates a new `StableMir` instance, with given test_function and arguments.
343            pub fn new(args: Vec<String>, callback: F) -> Self {
344                StableMir { args, callback: Some(callback), result: None }
345            }
346
347            /// Runs the compiler against given target and tests it with `test_function`
348            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                    // Two cases here:
360                    // - `run` finished normally and returned `Err`
361                    // - `run` panicked with `FatalErr`
362                    // You might think that normal compile errors cause the former, and
363                    // ICEs cause the latter. But some normal compiler errors also cause
364                    // the latter. So we can't meaningfully distinguish them, and group
365                    // them together.
366                    (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            /// Called after analysis. Return value instructs the compiler whether to
378            /// continue the compilation afterwards (defaults to `Compilation::Continue`)
379            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
404/// Similar to rustc's `FxIndexMap`, `IndexMap` with extra
405/// safety features added.
406pub 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
436/// Trait used to translate a stable construct to its rustc counterpart.
437///
438/// This is basically a mirror of [crate::rustc_smir::Stable].
439pub trait RustcInternal {
440    type T<'tcx>;
441    fn internal<'tcx>(&self, tables: &mut Tables<'_>, tcx: TyCtxt<'tcx>) -> Self::T<'tcx>;
442}