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 + Clone + 'static {
125        fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RLocalDefId, s: &S) -> 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        ) -> Self {
132            panic!("Can't convert MIR to the requested body type")
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, 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), s)
165    }
166
167    mod implementations {
168        use super::*;
169        impl IsBody for () {
170            fn body<'tcx, S: UnderOwnerState<'tcx>>(_did: RLocalDefId, _s: &S) -> Self {}
171            fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
172                _s: &S,
173                _body: rustc_middle::mir::Body<'tcx>,
174            ) -> Self {
175                ()
176            }
177        }
178        impl IsBody for ThirBody {
179            fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RLocalDefId, s: &S) -> Self {
180                let (thir, expr) = get_thir(did, s);
181                if *CORE_EXTRACTION_MODE {
182                    let expr = &thir.exprs[expr];
183                    Decorated {
184                        contents: Box::new(ExprKind::Tuple { fields: vec![] }),
185                        hir_id: None,
186                        attributes: vec![],
187                        ty: expr.ty.sinto(s),
188                        span: expr.span.sinto(s),
189                    }
190                } else {
191                    expr.sinto(&with_owner_id(s.base(), thir, (), did.to_def_id()))
192                }
193            }
194        }
195
196        impl<A: IsBody, B: IsBody> IsBody for (A, B) {
197            fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RLocalDefId, s: &S) -> Self {
198                (A::body(did, s), B::body(did, s))
199            }
200        }
201
202        impl<MirKind: IsMirKind + Clone + 'static> IsBody for MirBody<MirKind> {
203            fn body<'tcx, S: UnderOwnerState<'tcx>>(did: RLocalDefId, s: &S) -> Self {
204                let (thir, _) = get_thir(did, s);
205                let mir = MirKind::get_mir(s.base().tcx, did, |body| {
206                    let body = Rc::new(body.clone());
207                    body.sinto(&with_owner_id(
208                        s.base(),
209                        thir,
210                        body.clone(),
211                        did.to_def_id(),
212                    ))
213                });
214                mir.s_unwrap(s)
215            }
216            fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
217                s: &S,
218                body: rustc_middle::mir::Body<'tcx>,
219            ) -> Self {
220                let body = Rc::new(body.clone());
221                let s = &State {
222                    base: s.base(),
223                    owner_id: s.owner_id(),
224                    mir: body.clone(),
225                    binder: (),
226                    thir: (),
227                };
228                body.sinto(s)
229            }
230        }
231    }
232
233    impl<'tcx, S: UnderOwnerState<'tcx>, Body: IsBody> SInto<S, Body> for rustc_hir::BodyId {
234        fn sinto(&self, s: &S) -> Body {
235            body_from_id::<Body, _>(*self, s)
236        }
237    }
238}
hax_frontend_exporter/body.rs

hax_frontend_exporter/
body.rs
