hax_frontend_exporter/
state.rs

1use crate::prelude::*;
2use paste::paste;
3
4macro_rules! mk_aux {
5    ($state:ident {$($lts:lifetime)*} $field:ident {$($field_type:tt)+} {$($gen:tt)*} {$($gen_full:tt)*} {$($params:tt)*} {$($fields:tt)*}) => {
6        paste ! {
7            pub trait [<Has $field:camel>]<$($lts,)*> {
8                fn $field(self: &Self) -> $($field_type)+<$($lts)*>;
9            }
10            impl<$($lts,)*$($gen)*> [<Has $field:camel>]<$($lts,)*> for $state<$($params)*> {
11                fn $field(self: &Self) -> $($field_type)+<$($lts)*> {
12                    self.$field.clone()
13                }
14            }
15            pub trait [<Has $field:camel Setter>]<$($lts)*> {
16                type Out;
17                // fn [<with_ $field>]<$($lts,)*$($gen)*>(self: Self, $field: $($field_type)+<$($lts)*>) -> $state<$($params)*>;
18                fn [<with_ $field>](self: Self, $field: $($field_type)+<$($lts)*>) -> Self::Out;
19            }
20            // const _: &str = stringify!(
21            #[allow(unused)]
22            impl<$($lts,)*$($gen_full)*> [<Has $field:camel Setter>]<$($lts,)*> for $state<$($gen_full)*> {
23                type Out = $state<$($params)*>;
24                fn [<with_ $field>](self: Self, $field: $($field_type)+<$($lts)*>) -> Self::Out {
25                    let __this_field = $field;
26                    let $state { $($fields,)* } = self;
27                    let $field = __this_field;
28                    $state { $($fields,)* }
29                }
30            }
31            // );
32            // pub trait [<Has $field:camel Setter>]<$($lts,)*$($gen_full)*> {
33            //     fn [<with_ $field>](self: Self, $field: $($field_type)+<$($lts)*>) -> $state<$($params)*>;
34            // }
35            // impl<$($lts,)*$($gen_full)*> [<Has $field:camel Setter>]<$($lts,)*$($gen_full)*> for $state<$($gen_full)*> {
36            //     fn [<with_ $field>](self: Self, $field: $($field_type)+<$($lts)*>) -> $state<$($params)*> {
37            //         let __this_field = $field;
38            //         let $state { $($fields,)* } = self;
39            //         let $field = __this_field;
40            //         $state { $($fields,)* }
41            //     }
42            // }
43        }
44    };
45}
46macro_rules! mk_is_state_trait {
47    ($lts:tt {$($finished:tt)*} {{$f0:ident, $($l:tt)*} $($f:tt)*} $($generic:tt)*) => {
48        paste! {mk_is_state_trait!{
49                $lts {$($finished)* [<Has $f0:camel Setter>] <$($l)*> +} {$($f)*} $($generic)*
50                // $lts {$($finished)* [<Has $f0:camel Setter>] <$($l)* $($generic)*> +} {$($f)*} $($generic)*
51        }}
52    };
53    ({$($glts:lifetime,)*} {$($finished:tt)*} {} $($generic:tt)*) => {
54        pub trait IsState<$($glts,)*> = $($finished)*;
55    };
56}
57macro_rules! mk {
58    (struct $state:ident<$($glts:lifetime),*> {$($field:ident : {$($lts:lifetime),*} $field_type:ty),*$(,)?}) => {
59        mk!(@$state {} {$($field)*} {$($field: {$($lts),*} {$field_type},)*});
60        mk_is_state_trait!({$($glts,)*} {} {$({$field, $($lts,)*})*} $([<$field:camel>],)*);
61    };
62    (@$state:ident {$($acc:tt)*} $fields:tt {
63        $field:ident : $lts:tt $field_type:tt
64        $(,$($rest:tt)*)?
65    }) => {mk!(@$state {
66        $($acc)* $fields $field: $lts $field_type,
67    } $fields {$($($rest)*)?} );};
68    (@$state:ident $body:tt $fields:tt {$(,)?}) => { mk! (@@ $state $body ); };
69    (@@$state:ident {$({$($fields:tt)*} $field:ident : {$($lts:lifetime)*} {$($field_type:tt)+},)*}) => {
70        paste! {
71            #[derive(Clone)]
72            pub struct $state<$([<$field:camel>],)*>{
73                $(pub $field: [<$field:camel>],)*
74            }
75        }
76        $(
77            macro_rules! __inner_helper {
78                ($gen:tt {$$($full_gen:tt)*} {$$($params:tt)*} $field $$($rest:tt)*) => {
79                    paste! {__inner_helper!(
80                        $gen {$$($full_gen)*[<$field:camel>],}
81                        {$$($params)*$($field_type)+<$($lts,)*>,} $$($rest)*
82                    );}
83                };
84                ({$$($gen:tt)*} {$$($full_gen:tt)*} {$$($params:tt)*} $i:ident $$($rest:tt)*) => {
85                    paste! {__inner_helper!(
86                        {$$($gen)*[<$i:camel>],} {$$($full_gen)*[<$i:camel>],}
87                        {$$($params)*[<$i:camel>],} $$($rest)*
88                    );}
89                };
90                ($gen:tt $full_gen:tt $params:tt $$(,)?) => {
91                    mk_aux!($state {$($lts)*} $field {$($field_type)+} $gen $full_gen $params {$($fields)*});
92                };
93            }
94            __inner_helper!({} {} {} $($fields)*);
95        )*
96    };
97}
98
99mod types {
100    use crate::prelude::*;
101    use rustc_middle::ty;
102    use std::{cell::RefCell, sync::Arc};
103
104    pub struct LocalContextS {
105        pub vars: HashMap<rustc_middle::thir::LocalVarId, String>,
106    }
107
108    impl Default for LocalContextS {
109        fn default() -> Self {
110            Self::new()
111        }
112    }
113
114    impl LocalContextS {
115        pub fn new() -> LocalContextS {
116            LocalContextS {
117                vars: HashMap::new(),
118            }
119        }
120    }
121
122    /// Global caches
123    #[derive(Default)]
124    pub struct GlobalCache<'tcx> {
125        /// Cache the `Span` translations.
126        pub spans: HashMap<rustc_span::Span, Span>,
127        /// Per-item cache.
128        pub per_item: HashMap<RDefId, ItemCache<'tcx>>,
129        /// A ID table session, providing fresh IDs.
130        pub id_table_session: id_table::Session,
131    }
132
133    /// Defines a mapping from types to types, for use with `TypeMap`.
134    pub struct FullDefMapper {}
135    impl TypeMapper for FullDefMapper {
136        type Value<Body: TypeMappable> = Arc<FullDef<Body>>;
137    }
138
139    /// Defines a mapping from types to types, for use with `TypeMap`.
140    pub struct PromotedFullDefsMapper {}
141    impl TypeMapper for PromotedFullDefsMapper {
142        type Value<Body: TypeMappable> = HashMap<PromotedId, Arc<FullDef<Body>>>;
143    }
144
145    /// Per-item cache
146    #[derive(Default)]
147    pub struct ItemCache<'tcx> {
148        /// The translated `DefId`.
149        pub def_id: Option<DefId>,
150        /// The translated definitions, generic in the Body kind.
151        pub full_def: TypeMap<FullDefMapper>,
152        /// The Promoted constants of this body, if any.
153        pub promoteds: TypeMap<PromotedFullDefsMapper>,
154        /// Cache the `Ty` translations.
155        pub tys: HashMap<ty::Ty<'tcx>, Ty>,
156        /// Cache the `ItemRef` translations.
157        pub item_refs: HashMap<(RDefId, ty::GenericArgsRef<'tcx>), ItemRef>,
158        /// Cache the trait resolution engine for each item.
159        pub predicate_searcher: Option<crate::traits::PredicateSearcher<'tcx>>,
160        /// Cache of trait refs to resolved impl expressions.
161        pub impl_exprs: HashMap<ty::PolyTraitRef<'tcx>, crate::traits::ImplExpr>,
162    }
163
164    #[derive(Clone)]
165    pub struct Base<'tcx> {
166        pub options: Rc<hax_frontend_exporter_options::Options>,
167        pub local_ctx: Rc<RefCell<LocalContextS>>,
168        pub opt_def_id: Option<rustc_hir::def_id::DefId>,
169        pub cache: Rc<RefCell<GlobalCache<'tcx>>>,
170        pub tcx: ty::TyCtxt<'tcx>,
171        /// Rust doesn't enforce bounds on generic parameters in type
172        /// aliases. Thus, when translating type aliases, we need to
173        /// disable the resolution of implementation expressions. For
174        /// more details, please see
175        /// https://github.com/hacspec/hax/issues/707.
176        pub ty_alias_mode: bool,
177    }
178
179    impl<'tcx> Base<'tcx> {
180        pub fn new(
181            tcx: rustc_middle::ty::TyCtxt<'tcx>,
182            options: hax_frontend_exporter_options::Options,
183        ) -> Self {
184            Self {
185                tcx,
186                cache: Default::default(),
187                options: Rc::new(options),
188                // Always prefer `s.owner_id()` to `s.base().opt_def_id`.
189                // `opt_def_id` is used in `utils` for error reporting
190                opt_def_id: None,
191                local_ctx: Rc::new(RefCell::new(LocalContextS::new())),
192                ty_alias_mode: false,
193            }
194        }
195    }
196
197    pub type MacroCalls = Rc<HashMap<Span, Span>>;
198    pub type RcThir<'tcx> = Rc<rustc_middle::thir::Thir<'tcx>>;
199    pub type RcMir<'tcx> = Rc<rustc_middle::mir::Body<'tcx>>;
200    pub type UnitBinder<'tcx> = rustc_middle::ty::Binder<'tcx, ()>;
201}
202
203mk!(
204    struct State<'tcx> {
205        base: {'tcx} types::Base,
206        thir: {'tcx} types::RcThir,
207        mir: {'tcx} types::RcMir,
208        owner_id: {} rustc_hir::def_id::DefId,
209        binder: {'tcx} types::UnitBinder,
210    }
211);
212
213pub use self::types::*;
214
215pub type StateWithBase<'tcx> = State<Base<'tcx>, (), (), (), ()>;
216pub type StateWithOwner<'tcx> = State<Base<'tcx>, (), (), rustc_hir::def_id::DefId, ()>;
217pub type StateWithBinder<'tcx> =
218    State<Base<'tcx>, (), (), rustc_hir::def_id::DefId, types::UnitBinder<'tcx>>;
219pub type StateWithThir<'tcx> =
220    State<Base<'tcx>, types::RcThir<'tcx>, (), rustc_hir::def_id::DefId, ()>;
221pub type StateWithMir<'tcx> =
222    State<Base<'tcx>, (), types::RcMir<'tcx>, rustc_hir::def_id::DefId, ()>;
223
224impl<'tcx> StateWithBase<'tcx> {
225    pub fn new(
226        tcx: rustc_middle::ty::TyCtxt<'tcx>,
227        options: hax_frontend_exporter_options::Options,
228    ) -> Self {
229        Self {
230            thir: (),
231            mir: (),
232            owner_id: (),
233            binder: (),
234            base: Base::new(tcx, options),
235        }
236    }
237}
238
239impl<'tcx> StateWithOwner<'tcx> {
240    pub fn new_from_state_and_id<S: BaseState<'tcx>>(s: &S, id: rustc_hir::def_id::DefId) -> Self {
241        State {
242            thir: (),
243            mir: (),
244            owner_id: id,
245            binder: (),
246            base: s.base().clone(),
247        }
248    }
249}
250impl<'tcx> StateWithMir<'tcx> {
251    pub fn new_from_mir(
252        tcx: rustc_middle::ty::TyCtxt<'tcx>,
253        options: hax_frontend_exporter_options::Options,
254        mir: rustc_middle::mir::Body<'tcx>,
255        owner_id: rustc_hir::def_id::DefId,
256    ) -> Self {
257        Self {
258            thir: (),
259            mir: Rc::new(mir),
260            owner_id,
261            binder: (),
262            base: Base::new(tcx, options),
263        }
264    }
265}
266impl<'tcx> StateWithThir<'tcx> {
267    pub fn from_thir(
268        base: Base<'tcx>,
269        owner_id: rustc_hir::def_id::DefId,
270        thir: types::RcThir<'tcx>,
271    ) -> Self {
272        Self {
273            thir,
274            mir: (),
275            owner_id,
276            binder: (),
277            base,
278        }
279    }
280}
281
282impl<'tcx> StateWithOwner<'tcx> {
283    pub fn from_under_owner<S: UnderOwnerState<'tcx>>(s: &S) -> Self {
284        Self {
285            base: s.base(),
286            owner_id: s.owner_id(),
287            thir: (),
288            mir: (),
289            binder: (),
290        }
291    }
292}
293
294/// Updates the OnwerId in a state, making sure to override `opt_def_id` in base as well.
295pub fn with_owner_id<THIR, MIR>(
296    mut base: types::Base<'_>,
297    thir: THIR,
298    mir: MIR,
299    owner_id: rustc_hir::def_id::DefId,
300) -> State<types::Base<'_>, THIR, MIR, rustc_hir::def_id::DefId, ()> {
301    base.opt_def_id = Some(owner_id);
302    State {
303        thir,
304        owner_id,
305        base,
306        mir,
307        binder: (),
308    }
309}
310
311pub trait BaseState<'tcx> = HasBase<'tcx> + Clone + IsState<'tcx>;
312
313/// State of anything below a `owner_id`.
314pub trait UnderOwnerState<'tcx> = BaseState<'tcx> + HasOwnerId;
315
316/// State of anything below a binder.
317pub trait UnderBinderState<'tcx> = UnderOwnerState<'tcx> + HasBinder<'tcx>;
318
319/// While translating expressions, we expect to always have a THIR
320/// body and an `owner_id` in the state
321pub trait ExprState<'tcx> = UnderOwnerState<'tcx> + HasThir<'tcx>;
322
323pub trait WithGlobalCacheExt<'tcx>: BaseState<'tcx> {
324    /// Access the global cache. You must not call `sinto` within this function as this will likely
325    /// result in `BorrowMut` panics.
326    fn with_global_cache<T>(&self, f: impl FnOnce(&mut GlobalCache<'tcx>) -> T) -> T {
327        let base = self.base();
328        let mut cache = base.cache.borrow_mut();
329        f(&mut *cache)
330    }
331    /// Access the cache for a given item. You must not call `sinto` within this function as this
332    /// will likely result in `BorrowMut` panics.
333    fn with_item_cache<T>(&self, def_id: RDefId, f: impl FnOnce(&mut ItemCache<'tcx>) -> T) -> T {
334        self.with_global_cache(|cache| f(cache.per_item.entry(def_id).or_default()))
335    }
336}
337impl<'tcx, S: BaseState<'tcx>> WithGlobalCacheExt<'tcx> for S {}
338
339pub trait WithItemCacheExt<'tcx>: UnderOwnerState<'tcx> {
340    /// Access the cache for the current item. You must not call `sinto` within this function as
341    /// this will likely result in `BorrowMut` panics.
342    fn with_cache<T>(&self, f: impl FnOnce(&mut ItemCache<'tcx>) -> T) -> T {
343        self.with_item_cache(self.owner_id(), f)
344    }
345}
346impl<'tcx, S: UnderOwnerState<'tcx>> WithItemCacheExt<'tcx> for S {}
347
348impl ImplInfos {
349    fn from(base: Base<'_>, did: rustc_hir::def_id::DefId) -> Self {
350        let tcx = base.tcx;
351        let s = &with_owner_id(base, (), (), did);
352
353        Self {
354            generics: tcx.generics_of(did).sinto(s),
355            typ: tcx.type_of(did).instantiate_identity().sinto(s),
356            trait_ref: tcx
357                .impl_trait_ref(did)
358                .map(|trait_ref| trait_ref.instantiate_identity())
359                .sinto(s),
360            clauses: predicates_defined_on(tcx, did).predicates.sinto(s),
361        }
362    }
363}
364
365/// Returns a map from every implementation (`Impl`) `DefId`s to the
366/// type they implement, plus the bounds.
367pub fn impl_def_ids_to_impled_types_and_bounds<'tcx, S: BaseState<'tcx>>(
368    s: &S,
369) -> HashMap<DefId, ImplInfos> {
370    let tcx = s.base().tcx;
371
372    let def_ids: Vec<_> = s.with_global_cache(|cache| cache.per_item.keys().copied().collect());
373    let with_parents = |mut did: rustc_hir::def_id::DefId| {
374        let mut acc = vec![did];
375        while let Some(parent) = tcx.opt_parent(did) {
376            did = parent;
377            acc.push(did);
378        }
379        acc.into_iter()
380    };
381    use itertools::Itertools;
382    def_ids
383        .into_iter()
384        .flat_map(with_parents)
385        .unique()
386        .filter(|&did| {
387            // keep only DefIds that corresponds to implementations
388            matches!(
389                tcx.def_path(did).data.last(),
390                Some(rustc_hir::definitions::DisambiguatedDefPathData {
391                    data: rustc_hir::definitions::DefPathData::Impl,
392                    ..
393                })
394            )
395        })
396        .map(|did| (did.sinto(s), ImplInfos::from(s.base(), did)))
397        .collect()
398}