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

hax_frontend_exporter/
body.rs
