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 #[derive(Default)]
90 pub struct GlobalCache<'tcx> {
91 pub per_item: HashMap<DefId, ItemCache<'tcx>>,
93 pub def_ids: HashMap<RDefId, DefId>,
95 pub reverse_item_refs_map: HashMap<ItemRef, ty::GenericArgsRef<'tcx>>,
97 pub synthetic_item_data: HashMap<SyntheticItem, SyntheticItemData<'tcx>>,
99 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 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 #[derive(Default)]
152 pub struct ItemCache<'tcx> {
153 pub full_defs:
158 HashMap<(Option<PromotedId>, Option<ty::GenericArgsRef<'tcx>>), Arc<FullDef<'tcx>>>,
159 pub tys: HashMap<ty::Ty<'tcx>, Ty>,
161 pub item_refs: HashMap<(DefId, ty::GenericArgsRef<'tcx>, AssocItemResolution), ItemRef>,
163 pub trait_proofs: HashMap<ty::PolyTraitRef<'tcx>, crate::hax::traits::TraitProof>,
165 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 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 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 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 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
261pub 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
286pub trait UnderBinderState<'tcx> = UnderOwnerState<'tcx> + HasBinder<'tcx>;
288
289pub trait WithGlobalCacheExt<'tcx>: BaseState<'tcx> {
290 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 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 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 {}