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: Sized + std::fmt::Debug + Clone + 'static {
126        fn body<'tcx, S: UnderOwnerState<'tcx>>(
127            s: &S,
128            did: RDefId,
129            instantiate: Option<ty::GenericArgsRef<'tcx>>,
130        ) -> Option<Self>;
131
132        /// Reuse a MIR body we already got. Panic if that's impossible.
133        fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
134            _s: &S,
135            _body: rustc_middle::mir::Body<'tcx>,
136        ) -> Option<Self> {
137            None
138        }
139    }
140
141    pub fn make_fn_def<'tcx, Body: IsBody, S: BaseState<'tcx>>(
142        fn_sig: &rustc_hir::FnSig,
143        body_id: &rustc_hir::BodyId,
144        s: &S,
145    ) -> FnDef<Body> {
146        let hir_id = body_id.hir_id;
147        let ldid = hir_id.owner.def_id;
148
149        let (thir, expr_entrypoint) = get_thir(ldid, s);
150        let s = &s.with_owner_id(ldid.to_def_id()).with_thir(thir.clone());
151        FnDef {
152            params: thir.params.raw.sinto(s),
153            ret: thir.exprs[expr_entrypoint].ty.sinto(s),
154            body: Body::body(s, ldid.to_def_id(), None).s_unwrap(s),
155            sig_span: fn_sig.span.sinto(s),
156            header: fn_sig.header.sinto(s),
157        }
158    }
159
160    pub fn body_from_id<'tcx, Body: IsBody, S: UnderOwnerState<'tcx>>(
161        id: rustc_hir::BodyId,
162        s: &S,
163    ) -> Body {
164        // **Important:**
165        // We need a local id here, and we get it from the owner id, which must
166        // be local. It is safe to do so, because if we have access to HIR objects,
167        // it necessarily means we are exploring a local item (we don't have
168        // access to the HIR of external objects, only their MIR).
169        Body::body(s, s.base().tcx.hir_body_owner_def_id(id).to_def_id(), None).s_unwrap(s)
170    }
171
172    mod implementations {
173        use super::*;
174        impl IsBody for () {
175            fn body<'tcx, S: UnderOwnerState<'tcx>>(
176                _s: &S,
177                _did: RDefId,
178                _instantiate: Option<ty::GenericArgsRef<'tcx>>,
179            ) -> Option<Self> {
180                Some(())
181            }
182            fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
183                _s: &S,
184                _body: rustc_middle::mir::Body<'tcx>,
185            ) -> Option<Self> {
186                Some(())
187            }
188        }
189        impl IsBody for ThirBody {
190            fn body<'tcx, S: BaseState<'tcx>>(
191                s: &S,
192                did: RDefId,
193                instantiate: Option<ty::GenericArgsRef<'tcx>>,
194            ) -> Option<Self> {
195                let did = did.as_local()?;
196                let (thir, expr) = get_thir(did, s);
197                assert!(instantiate.is_none(), "monomorphized thir isn't supported");
198                let s = &s.with_owner_id(did.to_def_id());
199                Some(if *CORE_EXTRACTION_MODE {
200                    let expr = &thir.exprs[expr];
201                    Decorated {
202                        contents: Box::new(ExprKind::Tuple { fields: vec![] }),
203                        hir_id: None,
204                        attributes: vec![],
205                        ty: expr.ty.sinto(s),
206                        span: expr.span.sinto(s),
207                    }
208                } else {
209                    expr.sinto(&s.with_thir(thir))
210                })
211            }
212        }
213
214        impl<A: IsBody, B: IsBody> IsBody for (A, B) {
215            fn body<'tcx, S: UnderOwnerState<'tcx>>(
216                s: &S,
217                did: RDefId,
218                instantiate: Option<ty::GenericArgsRef<'tcx>>,
219            ) -> Option<Self> {
220                Some((A::body(s, did, instantiate)?, B::body(s, did, instantiate)?))
221            }
222        }
223
224        impl<MirKind: IsMirKind + Clone + 'static> IsBody for MirBody<MirKind> {
225            fn body<'tcx, S: UnderOwnerState<'tcx>>(
226                s: &S,
227                did: RDefId,
228                instantiate: Option<ty::GenericArgsRef<'tcx>>,
229            ) -> Option<Self> {
230                let tcx = s.base().tcx;
231                let typing_env = s.typing_env();
232                MirKind::get_mir(tcx, did, |body| {
233                    let body = substitute(tcx, typing_env, instantiate, body.clone());
234                    let body = Rc::new(body);
235                    body.sinto(&s.with_mir(body.clone()))
236                })
237            }
238            fn from_mir<'tcx, S: UnderOwnerState<'tcx>>(
239                s: &S,
240                body: rustc_middle::mir::Body<'tcx>,
241            ) -> Option<Self> {
242                let body = Rc::new(body.clone());
243                let s = &s.with_mir(body.clone());
244                Some(body.sinto(s))
245            }
246        }
247    }
248
249    impl<'tcx, S: UnderOwnerState<'tcx>, Body: IsBody> SInto<S, Body> for rustc_hir::BodyId {
250        fn sinto(&self, s: &S) -> Body {
251            body_from_id::<Body, _>(*self, s)
252        }
253    }
254}
hax_frontend_exporter/body.rs

hax_frontend_exporter/
body.rs
