hax_frontend_exporter/
body.rs

1pub use module::*;
2
3#[cfg(not(feature = "rustc"))]
4mod module {
5    pub trait IsBody: Sized + Clone + 'static {}
6    impl<T: Sized + Clone + 'static> IsBody for T {}
7}
8
9#[cfg(feature = "rustc")]
10mod module {
11    pub use crate::prelude::*;
12    pub use rustc_hir::{
13        def_id::{DefId as RDefId, LocalDefId as RLocalDefId},
14        hir_id::OwnerId as ROwnerId,
15    };
16
17    mod store {
18        //! This module helps at store bodies to avoid stealing.
19        //! `rustc_data_structures::steal::Steal` is a box for which the content can be stolen, for performance reasons.
20        //! The query system of Rust creates and steal such boxes, resulting in hax trying to borrow the value of a Steal while some query stole it already.
21        //! This module provides an ad-hoc global cache and query overrides to deal with this issue.
22        use rustc_hir::def_id::LocalDefId;
23        use rustc_middle::mir::Body;
24        use rustc_middle::query::plumbing::IntoQueryParam;
25        use rustc_middle::thir::{ExprId, Thir};
26        use std::cell::RefCell;
27        use std::collections::HashMap;
28        use std::rc::Rc;
29
30        thread_local! {
31            static THIR_BODY: RefCell<HashMap<LocalDefId, (Rc<Thir<'static>>, ExprId)>> = RefCell::new(HashMap::new());
32            static MIR_BUILT: RefCell<HashMap<LocalDefId, Rc<Body<'static>>>> = RefCell::new(HashMap::new());
33        }
34
35        /// Register overrides for rustc queries.
36        /// This will clone and store bodies for THIR and MIR (built) in an ad-hoc global cache.
37        pub fn override_queries_store_body(providers: &mut rustc_middle::query::Providers) {
38            providers.thir_body = |tcx, def_id| {
39                let (steal, expr_id) =
40                    (rustc_interface::DEFAULT_QUERY_PROVIDERS.thir_body)(tcx, def_id)?;
41                let body = steal.borrow().clone();
42                let body: Thir<'static> = unsafe { std::mem::transmute(body) };
43                THIR_BODY.with(|map| map.borrow_mut().insert(def_id, (Rc::new(body), expr_id)));
44                Ok((steal, expr_id))
45            };
46            providers.mir_built = |tcx, def_id| {
47                let steal = (rustc_interface::DEFAULT_QUERY_PROVIDERS.mir_built)(tcx, def_id);
48                let body = steal.borrow().clone();
49                let body: Body<'static> = unsafe { std::mem::transmute(body) };
50                MIR_BUILT.with(|map| map.borrow_mut().insert(def_id, Rc::new(body)));
51                steal
52            };
53        }
54
55        /// Extension trait that provides non-stealing variants of `thir_body` and `mir_built`.
56        /// Those methods requires rustc queries to be overriden with the helper function `register` above.
57        #[extension_traits::extension(pub trait SafeTyCtxtBodies)]
58        impl<'tcx> rustc_middle::ty::TyCtxt<'tcx> {
59            fn thir_body_safe(
60                &self,
61                key: impl IntoQueryParam<rustc_span::def_id::LocalDefId>,
62            ) -> Result<(Rc<Thir<'tcx>>, ExprId), rustc_span::ErrorGuaranteed> {
63                let key = key.into_query_param();
64                if !THIR_BODY.with(|map| map.borrow().contains_key(&key)) {
65                    // Compute a body, which will insert a body in `THIR_BODIES`.
66                    let _ = self.thir_body(key);
67                }
68                THIR_BODY.with(|map| {
69                    let (body, expr) = map
70                        .borrow_mut()
71                        .get(&key)
72                        .expect("Did we forgot to call `register`?")
73                        .clone();
74                    let body: Rc<Thir<'tcx>> = unsafe { std::mem::transmute(body) };
75                    Ok((body, expr))
76                })
77            }
78            fn mir_built_safe(
79                &self,
80                key: impl IntoQueryParam<rustc_span::def_id::LocalDefId>,
81            ) -> Rc<Body<'tcx>> {
82                let key = key.into_query_param();
83                if !MIR_BUILT.with(|map| map.borrow().contains_key(&key)) {
84                    // Compute a body, which will insert a body in `MIR_BODIES`.
85                    let _ = self.mir_built(key);
86                }
87                MIR_BUILT.with(|map| {
88                    let body = map
89                        .borrow_mut()
90                        .get(&key)
91                        .expect("Did we forgot to call `register`?")
92                        .clone();
93                    unsafe { std::mem::transmute(body) }
94                })
95            }
96        }
97    }
98    pub use store::*;
99
100    pub fn get_thir<'tcx, S: UnderOwnerState<'tcx>>(
101        did: RLocalDefId,
102        s: &S,
103    ) -> (
104        Rc<rustc_middle::thir::Thir<'tcx>>,
105        rustc_middle::thir::ExprId,
106    ) {
107        let tcx = s.base().tcx;
108
109        // The `type_of` anon constants isn't available directly, it needs to be fed by some
110        // other query. This hack ensures this happens, otherwise `thir_body` returns an error.
111        // See https://rust-lang.zulipchat.com/#narrow/channel/182449-t-compiler.2Fhelp/topic/Change.20in.20THIR.20of.20anonymous.20constants.3F/near/509764021 .
112        let hir_id = tcx.local_def_id_to_hir_id(did);
113        for (parent_id, parent) in tcx.hir_parent_iter(hir_id) {
114            if let rustc_hir::Node::Item(..) = parent {
115                let _ = tcx.check_well_formed(parent_id.owner.def_id);
116                break;
117            }
118        }
119
120        let msg = |_| fatal!(s[tcx.def_span(did)], "THIR not found for {:?}", did);
121        tcx.thir_body_safe(did).as_ref().unwrap_or_else(msg).clone()
122    }
123
124    pub trait IsBody: Sized + std::fmt::Debug + Clone + 'static {
125        fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RDefId, s: &S) -> Option<Self>;
126
127        /// Reuse a MIR body we already got. Panic if that's impossible.
128        fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
129            _s: &S,
130            _body: rustc_middle::mir::Body<'tcx>,
131        ) -> Option<Self> {
132            None
133        }
134    }
135
136    pub fn make_fn_def<'tcx, Body: IsBody, S: UnderOwnerState<'tcx>>(
137        fn_sig: &rustc_hir::FnSig,
138        body_id: &rustc_hir::BodyId,
139        s: &S,
140    ) -> FnDef<Body> {
141        let hir_id = body_id.hir_id;
142        let ldid = hir_id.owner.def_id;
143
144        let (thir, expr_entrypoint) = get_thir(ldid, s);
145        let s = &with_owner_id(s.base(), thir.clone(), (), ldid.to_def_id());
146        FnDef {
147            params: thir.params.raw.sinto(s),
148            ret: thir.exprs[expr_entrypoint].ty.sinto(s),
149            body: Body::body(ldid.to_def_id(), s).s_unwrap(s),
150            sig_span: fn_sig.span.sinto(s),
151            header: fn_sig.header.sinto(s),
152        }
153    }
154
155    pub fn body_from_id<'tcx, Body: IsBody, S: UnderOwnerState<'tcx>>(
156        id: rustc_hir::BodyId,
157        s: &S,
158    ) -> Body {
159        // **Important:**
160        // We need a local id here, and we get it from the owner id, which must
161        // be local. It is safe to do so, because if we have access to HIR objects,
162        // it necessarily means we are exploring a local item (we don't have
163        // access to the HIR of external objects, only their MIR).
164        Body::body(s.base().tcx.hir_body_owner_def_id(id).to_def_id(), s).s_unwrap(s)
165    }
166
167    mod implementations {
168        use super::*;
169        impl IsBody for () {
170            fn body<'tcx, S: UnderOwnerState<'tcx>>(_did: RDefId, _s: &S) -> Option<Self> {
171                Some(())
172            }
173            fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
174                _s: &S,
175                _body: rustc_middle::mir::Body<'tcx>,
176            ) -> Option<Self> {
177                Some(())
178            }
179        }
180        impl IsBody for ThirBody {
181            fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RDefId, s: &S) -> Option<Self> {
182                let did = did.as_local()?;
183                let (thir, expr) = get_thir(did, s);
184                Some(if *CORE_EXTRACTION_MODE {
185                    let expr = &thir.exprs[expr];
186                    Decorated {
187                        contents: Box::new(ExprKind::Tuple { fields: vec![] }),
188                        hir_id: None,
189                        attributes: vec![],
190                        ty: expr.ty.sinto(s),
191                        span: expr.span.sinto(s),
192                    }
193                } else {
194                    expr.sinto(&with_owner_id(s.base(), thir, (), did.to_def_id()))
195                })
196            }
197        }
198
199        impl<A: IsBody, B: IsBody> IsBody for (A, B) {
200            fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RDefId, s: &S) -> Option<Self> {
201                Some((A::body(did, s)?, B::body(did, s)?))
202            }
203        }
204
205        impl<MirKind: IsMirKind + Clone + 'static> IsBody for MirBody<MirKind> {
206            fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RDefId, s: &S) -> Option<Self> {
207                MirKind::get_mir(s.base().tcx, did, |body| {
208                    let body = Rc::new(body.clone());
209                    body.sinto(&with_owner_id(s.base(), (), body.clone(), did))
210                })
211            }
212            fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
213                s: &S,
214                body: rustc_middle::mir::Body<'tcx>,
215            ) -> Option<Self> {
216                let body = Rc::new(body.clone());
217                let s = &State {
218                    base: s.base(),
219                    owner_id: s.owner_id(),
220                    mir: body.clone(),
221                    binder: (),
222                    thir: (),
223                };
224                Some(body.sinto(s))
225            }
226        }
227    }
228
229    impl<'tcx, S: UnderOwnerState<'tcx>, Body: IsBody> SInto<S, Body> for rustc_hir::BodyId {
230        fn sinto(&self, s: &S) -> Body {
231            body_from_id::<Body, _>(*self, s)
232        }
233    }
234}
hax_frontend_exporter/body.rs

hax_frontend_exporter/
body.rs
