Skip to main content

charon_driver/hax/
state.rs

1use crate::hax::prelude::*;
2use paste::paste;
3use rustc_middle::ty::{self, TyCtxt};
4
5macro_rules! mk_aux {
6    ($state:ident {$($lts:lifetime)*} $field:ident {$($field_type:tt)+} {$($gen:tt)*} {$($gen_full:tt)*} {$($params:tt)*} {$($fields:tt)*}) => {
7        paste ! {
8            pub trait [<Has $field:camel>]<$($lts,)*> {
9                fn $field(self: &Self) -> $($field_type)+<$($lts)*>;
10            }
11            impl<$($lts,)*$($gen)*> [<Has $field:camel>]<$($lts,)*> for $state<$($params)*> {
12                fn $field(self: &Self) -> $($field_type)+<$($lts)*> {
13                    self.$field.clone()
14                }
15            }
16        }
17    };
18}
19macro_rules! mk {
20    (struct $state:ident<$($glts:lifetime),*> {$($field:ident : {$($lts:lifetime),*} $field_type:ty),*$(,)?}) => {
21        mk!(@$state {} {$($field)*} {$($field: {$($lts),*} {$field_type},)*});
22    };
23    (@$state:ident {$($acc:tt)*} $fields:tt {
24        $field:ident : $lts:tt $field_type:tt
25        $(,$($rest:tt)*)?
26    }) => {mk!(@$state {
27        $($acc)* $fields $field: $lts $field_type,
28    } $fields {$($($rest)*)?} );};
29    (@$state:ident $body:tt $fields:tt {$(,)?}) => { mk! (@@ $state $body ); };
30    (@@$state:ident {$({$($fields:tt)*} $field:ident : {$($lts:lifetime)*} {$($field_type:tt)+},)*}) => {
31        paste! {
32            #[derive(Clone)]
33            pub struct $state<$([<$field:camel>],)*>{
34                $(pub $field: [<$field:camel>],)*
35            }
36        }
37        $(
38            macro_rules! __inner_helper {
39                ($gen:tt {$$($full_gen:tt)*} {$$($params:tt)*} $field $$($rest:tt)*) => {
40                    paste! {__inner_helper!(
41                        $gen {$$($full_gen)*[<$field:camel>],}
42                        {$$($params)*$($field_type)+<$($lts,)*>,} $$($rest)*
43                    );}
44                };
45                ({$$($gen:tt)*} {$$($full_gen:tt)*} {$$($params:tt)*} $i:ident $$($rest:tt)*) => {
46                    paste! {__inner_helper!(
47                        {$$($gen)*[<$i:camel>],} {$$($full_gen)*[<$i:camel>],}
48                        {$$($params)*[<$i:camel>],} $$($rest)*
49                    );}
50                };
51                ($gen:tt $full_gen:tt $params:tt $$(,)?) => {
52                    mk_aux!($state {$($lts)*} $field {$($field_type)+} $gen $full_gen $params {$($fields)*});
53                };
54            }
55            __inner_helper!({} {} {} $($fields)*);
56        )*
57    };
58}
59
60mod types {
61    use crate::hax::prelude::*;
62    use itertools::Itertools;
63    use rustc_hash::FxHashMap;
64    use rustc_hir::def_id::{CrateNum, LOCAL_CRATE};
65    use rustc_middle::ty;
66    use rustc_span::symbol::Symbol;
67    use rustc_trait_elaboration::ElaborationCtx;
68    use std::{cell::RefCell, sync::Arc};
69
70    pub struct LocalContextS {
71        pub vars: HashMap<rustc_middle::thir::LocalVarId, String>,
72    }
73
74    impl Default for LocalContextS {
75        fn default() -> Self {
76            Self::new()
77        }
78    }
79
80    impl LocalContextS {
81        pub fn new() -> LocalContextS {
82            LocalContextS {
83                vars: HashMap::new(),
84            }
85        }
86    }
87
88    /// Global caches
89    #[derive(Default)]
90    pub struct GlobalCache<'tcx> {
91        /// Per-item cache.
92        pub per_item: HashMap<DefId, ItemCache<'tcx>>,
93        /// Map rustc def ids to their hax counterpart.
94        pub def_ids: HashMap<RDefId, DefId>,
95        /// Map that recovers rustc args for a given `ItemRef`.
96        pub reverse_item_refs_map: HashMap<ItemRef, ty::GenericArgsRef<'tcx>>,
97        /// Data for synthetic items. See the `synthetic_items` module.
98        pub synthetic_item_data: HashMap<SyntheticItem, SyntheticItemData<'tcx>>,
99        /// Cached names and disambiguators for crate names.
100        pub disambiguated_crate_names: Option<FxHashMap<CrateNum, (Symbol, u32)>>,
101    }
102
103    impl<'tcx> GlobalCache<'tcx> {
104        pub fn crate_name(&mut self, tcx: ty::TyCtxt<'tcx>, krate: CrateNum) -> (Symbol, u32) {
105            self.disambiguated_crate_names
106                .get_or_insert_with(|| {
107                    std::iter::once(LOCAL_CRATE)
108                        .chain(tcx.crates(()).iter().copied())
109                        .into_group_map_by(|&cnum| tcx.crate_name(cnum))
110                        .into_iter()
111                        .filter(|(_, crates_by_this_name)| crates_by_this_name.len() > 1)
112                        .flat_map(|(name, mut crates_by_this_name)| {
113                            crates_by_this_name.sort_by_key(|&cnum| {
114                                // `tcx.stable_crate_id` isn't actually so stable across machines.
115                                // We try our own stability here based on paths.
116                                let source_file = tcx
117                                    .sess
118                                    .source_map()
119                                    .span_to_filename(tcx.def_span(cnum.as_def_id()))
120                                    .prefer_remapped_unconditionally()
121                                    .to_string_lossy()
122                                    .into_owned();
123                                let extern_paths = if cnum == LOCAL_CRATE {
124                                    Vec::new()
125                                } else {
126                                    tcx.crate_extern_paths(cnum)
127                                        .iter()
128                                        .map(|path| path.display().to_string())
129                                        .collect_vec()
130                                };
131                                (
132                                    cnum != LOCAL_CRATE,
133                                    source_file,
134                                    extern_paths,
135                                    cnum.as_u32(),
136                                )
137                            });
138                            crates_by_this_name.into_iter().enumerate().map(
139                                move |(disambiguator, cnum)| (cnum, (name, disambiguator as u32)),
140                            )
141                        })
142                        .collect()
143                })
144                .get(&krate)
145                .copied()
146                .unwrap_or_else(|| (tcx.crate_name(krate), 0))
147        }
148    }
149
150    /// Per-item cache
151    #[derive(Default)]
152    pub struct ItemCache<'tcx> {
153        /// The translated definitions, generic in the Body kind.
154        /// Each rustc `DefId` gives several hax `DefId`s: one for each promoted constant (if any),
155        /// and the base one represented by `None`. Moreover we can instantiate definitions with
156        /// generic arguments.
157        pub full_defs:
158            HashMap<(Option<PromotedId>, Option<ty::GenericArgsRef<'tcx>>), Arc<FullDef<'tcx>>>,
159        /// Cache the `Ty` translations.
160        pub tys: HashMap<ty::Ty<'tcx>, Ty>,
161        /// Cache the `ItemRef` translations. This is fast because `GenericArgsRef` is interned.
162        pub item_refs: HashMap<(DefId, ty::GenericArgsRef<'tcx>, AssocItemResolution), ItemRef>,
163        /// Cache of trait refs to resolved trait proofs.
164        pub trait_proofs: HashMap<ty::PolyTraitRef<'tcx>, crate::hax::traits::TraitProof>,
165        /// Generics for this item, if it is virtual.
166        pub virtual_generics: Option<&'tcx ty::Generics>,
167    }
168
169    #[derive(Clone)]
170    pub struct Base<'tcx> {
171        pub options: Rc<crate::hax::options::Options>,
172        pub local_ctx: Rc<RefCell<LocalContextS>>,
173        pub opt_def_id: Option<rustc_hir::def_id::DefId>,
174        pub cache: Rc<RefCell<GlobalCache<'tcx>>>,
175        pub elab_ctx: crate::hax::traits::ElaborationCtx<'tcx, DefId>,
176        pub tcx: ty::TyCtxt<'tcx>,
177    }
178
179    impl<'tcx> Base<'tcx> {
180        pub fn new(
181            tcx: rustc_middle::ty::TyCtxt<'tcx>,
182            options: crate::hax::options::Options,
183            bounds_options: crate::hax::options::BoundsOptions,
184        ) -> Self {
185            Self {
186                tcx,
187                cache: Default::default(),
188                options: Rc::new(options),
189                // Always prefer `s.owner_id()` to `s.base().opt_def_id`.
190                // `opt_def_id` is used in `utils` for error reporting
191                opt_def_id: None,
192                local_ctx: Rc::new(RefCell::new(LocalContextS::new())),
193                elab_ctx: ElaborationCtx::new(tcx, bounds_options),
194            }
195        }
196    }
197
198    pub type UnitBinder<'tcx> = rustc_middle::ty::Binder<'tcx, ()>;
199}
200
201mk!(
202    struct State<'tcx> {
203        base: {'tcx} types::Base,
204        owner: {} DefId,
205        binder: {'tcx} types::UnitBinder,
206    }
207);
208
209pub use self::types::*;
210
211pub type StateWithBase<'tcx> = State<Base<'tcx>, (), ()>;
212pub type StateWithOwner<'tcx> = State<Base<'tcx>, DefId, ()>;
213pub type StateWithBinder<'tcx> = State<Base<'tcx>, DefId, types::UnitBinder<'tcx>>;
214
215impl<'tcx> StateWithBase<'tcx> {
216    pub fn new(
217        tcx: rustc_middle::ty::TyCtxt<'tcx>,
218        options: crate::hax::options::Options,
219        bounds_options: crate::hax::options::BoundsOptions,
220    ) -> Self {
221        Self {
222            base: Base::new(tcx, options, bounds_options),
223            owner: (),
224            binder: (),
225        }
226    }
227}
228
229pub trait BaseState<'tcx>: HasBase<'tcx> + Clone {
230    fn tcx(&self) -> TyCtxt<'tcx> {
231        self.base().tcx
232    }
233
234    /// Create a state with the given owner.
235    fn with_hax_owner(&self, owner: &DefId) -> StateWithOwner<'tcx> {
236        let mut base = self.base();
237        base.opt_def_id = owner.as_real_or_promoted();
238        State {
239            owner: owner.clone(),
240            base,
241            binder: (),
242        }
243    }
244    /// Create a state with the given owner.
245    fn with_rustc_owner(&self, owner_id: RDefId) -> StateWithOwner<'tcx> {
246        let owner = &owner_id.sinto(self);
247        Self::with_hax_owner(self, owner)
248    }
249
250    /// State with only access to the global state.
251    fn base_state(&self) -> StateWithBase<'tcx> {
252        State {
253            base: self.base(),
254            owner: (),
255            binder: (),
256        }
257    }
258}
259impl<'tcx, T: HasBase<'tcx> + Clone> BaseState<'tcx> for T {}
260
261/// State of anything below an `owner`.
262pub trait UnderOwnerState<'tcx>: BaseState<'tcx> + HasOwner {
263    fn param_env(&self) -> ty::ParamEnv<'tcx> {
264        self.owner().param_env(self)
265    }
266    fn typing_env(&self) -> ty::TypingEnv<'tcx> {
267        self.owner().typing_env(self)
268    }
269    fn with_base(&self, base: types::Base<'tcx>) -> StateWithOwner<'tcx> {
270        State {
271            owner: self.owner().clone(),
272            base,
273            binder: (),
274        }
275    }
276    fn with_binder(&self, binder: types::UnitBinder<'tcx>) -> StateWithBinder<'tcx> {
277        State {
278            base: self.base(),
279            owner: self.owner().clone(),
280            binder,
281        }
282    }
283}
284impl<'tcx, T: BaseState<'tcx> + HasOwner> UnderOwnerState<'tcx> for T {}
285
286/// State of anything below a binder.
287pub trait UnderBinderState<'tcx> = UnderOwnerState<'tcx> + HasBinder<'tcx>;
288
289pub trait WithGlobalCacheExt<'tcx>: BaseState<'tcx> {
290    /// Access the global cache. You must not call `sinto` within this function as this will likely
291    /// result in `BorrowMut` panics.
292    fn with_global_cache<T>(&self, f: impl FnOnce(&mut GlobalCache<'tcx>) -> T) -> T {
293        let base = self.base();
294        let mut cache = base.cache.borrow_mut();
295        f(&mut cache)
296    }
297    /// Access the cache for a given item. You must not call `sinto` within this function as this
298    /// will likely result in `BorrowMut` panics.
299    fn with_item_cache<T>(&self, def_id: &DefId, f: impl FnOnce(&mut ItemCache<'tcx>) -> T) -> T {
300        self.with_global_cache(|cache| f(cache.per_item.entry(def_id.clone()).or_default()))
301    }
302}
303impl<'tcx, S: BaseState<'tcx>> WithGlobalCacheExt<'tcx> for S {}
304
305pub trait WithItemCacheExt<'tcx>: UnderOwnerState<'tcx> {
306    /// Access the cache for the current item. You must not call `sinto` within this function as
307    /// this will likely result in `BorrowMut` panics.
308    fn with_cache<T>(&self, f: impl FnOnce(&mut ItemCache<'tcx>) -> T) -> T {
309        self.with_item_cache(&self.owner(), f)
310    }
311    fn with_predicate_searcher<T>(
312        &self,
313        f: impl FnOnce(&mut PredicateSearcher<'tcx>, &StateWithBase<'tcx>) -> T,
314    ) -> T {
315        let s = self;
316        let base = s.base();
317        let owner = s.owner();
318        let base_state = s.base_state();
319        let mut predicate_searcher = base.elab_ctx.predicate_searcher_for(&base_state, owner);
320        f(&mut predicate_searcher, &base_state)
321    }
322}
323impl<'tcx, S: UnderOwnerState<'tcx>> WithItemCacheExt<'tcx> for S {}