flux_middle/
queries.rs

1use std::{
2    cell::{OnceCell, RefCell},
3    rc::Rc,
4};
5
6use flux_arc_interner::List;
7use flux_common::bug;
8use flux_errors::{E0999, ErrorGuaranteed};
9use flux_rustc_bridge::{
10    self, def_id_to_string,
11    lowering::{self, Lower, UnsupportedErr},
12    mir, ty,
13};
14use itertools::Itertools;
15use rustc_data_structures::unord::{ExtendUnord, UnordMap};
16use rustc_errors::Diagnostic;
17use rustc_hir::{
18    def::DefKind,
19    def_id::{CrateNum, DefId, LOCAL_CRATE, LocalDefId},
20};
21use rustc_index::IndexVec;
22use rustc_macros::{Decodable, Encodable};
23use rustc_span::{Span, Symbol};
24
25use crate::{
26    def_id::{FluxDefId, FluxId, FluxLocalDefId, MaybeExternId, ResolvedDefId},
27    fhir,
28    global_env::GlobalEnv,
29    rty::{
30        self,
31        refining::{self, Refine, Refiner},
32    },
33};
34
35type Cache<K, V> = RefCell<UnordMap<K, V>>;
36
37pub type QueryResult<T = ()> = Result<T, QueryErr>;
38
39/// An error produced by a query.
40///
41/// We make a distinction between errors reported at def-site and errors reported at use-site.
42///
43/// For most errors reported at the def-site of an item, it makes little sense to check the definition
44/// of dependent items. For example, if a function signature is ill-formed, checking the body of another
45/// function that calls it, can produce confusing errors. In some cases, we can even fail to produce
46/// a signature for a function in which case we can't even check its call sites. For these cases, we
47/// emit an error at the definition site and return a [`QueryErr::Emitted`]. When checking a dependent,
48/// we detect this and early return without reporting any errors at the use-site.
49///
50/// Other errors are better reported at the use-site. For example, if some code calls a function from
51/// an external crate that has unsupported features, we ought to report the error at the call-site,
52/// because it would be confusing to only mention the definition of the external function without
53/// showing which part of the code is calling it. To attach a span to an error one can use [`QueryErr::at`]
54/// to get a [`QueryErrAt`].
55///
56/// Both [`QueryErr`] and [`QueryErrAt`] implement [`Diagnostic`]. The implementation for [`QueryErr`]
57/// reports the error at the definition site, while the implementation for [`QueryErrAt`] reports it at
58/// the (attached) use-site span. This allows us to play a bit lose because we can emit an error without
59/// attaching a span, but this means we may forget to attach spans at some places. We should consider
60/// not implementing [`Diagnostic`] for [`QueryErr`] such that we always make the distinction between
61/// use-site and def-site explicit, e.g., we could have methods `QueryErr::at_use_site` and
62/// `QueryErr::at_def_site` returning types with different implementations of [`Diagnostic`].
63#[derive(Debug, Clone, Encodable, Decodable)]
64pub enum QueryErr {
65    Unsupported {
66        def_id: DefId,
67        err: UnsupportedErr,
68    },
69    Ignored {
70        def_id: DefId,
71    },
72    InvalidGenericArg {
73        def_id: DefId,
74    },
75    MissingAssocReft {
76        impl_id: DefId,
77        trait_id: DefId,
78        name: Symbol,
79    },
80    /// Used to report bugs, typically this means executing an arm in a match we thought it was
81    /// unreachable. Use this instead of panicking if it is easy to return a [`QueryErr`]. Use
82    /// [`QueryErr::bug`] or [`crate::query_bug!`] to construct this variant to track source location.
83    Bug {
84        def_id: Option<DefId>,
85        location: String,
86        msg: String,
87    },
88    Emitted(ErrorGuaranteed),
89}
90
91#[macro_export]
92macro_rules! query_bug {
93    ($fmt:literal $(,$args:expr)* $(,)?) => {
94        $crate::queries::QueryErr::bug(None, format_args!($fmt, $($args),*))
95    };
96    ($def_id:expr, $fmt:literal $(,$args:expr)* $(,)? ) => {{
97        $crate::queries::QueryErr::bug(Some($def_id.into()), format_args!($fmt, $($args),*))
98    }};
99}
100
101impl QueryErr {
102    pub fn unsupported(def_id: DefId, err: UnsupportedErr) -> Self {
103        QueryErr::Unsupported { def_id, err }
104    }
105
106    #[track_caller]
107    pub fn bug(def_id: Option<DefId>, msg: impl ToString) -> Self {
108        QueryErr::Bug {
109            def_id,
110            location: format!("{}", std::panic::Location::caller()),
111            msg: msg.to_string(),
112        }
113    }
114
115    pub fn at(self, span: Span) -> QueryErrAt {
116        QueryErrAt { span, err: self }
117    }
118}
119
120/// See [`QueryErr`]
121pub struct QueryErrAt {
122    span: Span,
123    err: QueryErr,
124}
125
126pub struct Providers {
127    pub collect_specs: fn(GlobalEnv) -> crate::Specs,
128    pub resolve_crate: fn(GlobalEnv) -> crate::ResolverOutput,
129    pub desugar: for<'genv> fn(
130        GlobalEnv<'genv, '_>,
131        LocalDefId,
132    ) -> QueryResult<UnordMap<LocalDefId, fhir::Node<'genv>>>,
133    pub fhir_crate: for<'genv> fn(GlobalEnv<'genv, '_>) -> fhir::FluxItems<'genv>,
134    pub qualifiers: fn(GlobalEnv) -> QueryResult<Vec<rty::Qualifier>>,
135    pub normalized_defns: fn(GlobalEnv) -> rty::NormalizedDefns,
136    pub func_sort: fn(GlobalEnv, FluxLocalDefId) -> rty::PolyFuncSort,
137    pub adt_sort_def_of: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::AdtSortDef>,
138    pub check_wf: for<'genv> fn(GlobalEnv, LocalDefId) -> QueryResult<Rc<rty::WfckResults>>,
139    pub adt_def: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::AdtDef>,
140    pub constant_info: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::ConstantInfo>,
141    pub type_of: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::EarlyBinder<rty::TyOrCtor>>,
142    pub variants_of: fn(
143        GlobalEnv,
144        LocalDefId,
145    ) -> QueryResult<rty::Opaqueness<rty::EarlyBinder<rty::PolyVariants>>>,
146    pub fn_sig: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::EarlyBinder<rty::PolyFnSig>>,
147    pub generics_of: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::Generics>,
148    pub refinement_generics_of:
149        fn(GlobalEnv, LocalDefId) -> QueryResult<rty::EarlyBinder<rty::RefinementGenerics>>,
150    pub predicates_of:
151        fn(GlobalEnv, LocalDefId) -> QueryResult<rty::EarlyBinder<rty::GenericPredicates>>,
152    pub assoc_refinements_of: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::AssocRefinements>,
153    pub sort_of_assoc_reft:
154        fn(GlobalEnv, FluxId<MaybeExternId>) -> QueryResult<rty::EarlyBinder<rty::FuncSort>>,
155    pub assoc_refinement_body:
156        fn(GlobalEnv, FluxId<MaybeExternId>) -> QueryResult<rty::EarlyBinder<rty::Lambda>>,
157    #[allow(clippy::type_complexity)]
158    pub default_assoc_refinement_body:
159        fn(GlobalEnv, FluxId<MaybeExternId>) -> QueryResult<Option<rty::EarlyBinder<rty::Lambda>>>,
160    pub item_bounds: fn(GlobalEnv, LocalDefId) -> QueryResult<rty::EarlyBinder<List<rty::Clause>>>,
161}
162
163macro_rules! empty_query {
164    () => {
165        flux_common::bug!("query not provided")
166    };
167}
168
169impl Default for Providers {
170    fn default() -> Self {
171        Self {
172            collect_specs: |_| empty_query!(),
173            resolve_crate: |_| empty_query!(),
174            desugar: |_, _| empty_query!(),
175            fhir_crate: |_| empty_query!(),
176            normalized_defns: |_| empty_query!(),
177            func_sort: |_, _| empty_query!(),
178            qualifiers: |_| empty_query!(),
179            adt_sort_def_of: |_, _| empty_query!(),
180            check_wf: |_, _| empty_query!(),
181            adt_def: |_, _| empty_query!(),
182            type_of: |_, _| empty_query!(),
183            variants_of: |_, _| empty_query!(),
184            fn_sig: |_, _| empty_query!(),
185            generics_of: |_, _| empty_query!(),
186            refinement_generics_of: |_, _| empty_query!(),
187            predicates_of: |_, _| empty_query!(),
188            assoc_refinements_of: |_, _| empty_query!(),
189            assoc_refinement_body: |_, _| empty_query!(),
190            default_assoc_refinement_body: |_, _| empty_query!(),
191            sort_of_assoc_reft: |_, _| empty_query!(),
192            item_bounds: |_, _| empty_query!(),
193            constant_info: |_, _| empty_query!(),
194        }
195    }
196}
197
198pub struct Queries<'genv, 'tcx> {
199    pub(crate) providers: Providers,
200    mir: Cache<LocalDefId, QueryResult<Rc<mir::Body<'tcx>>>>,
201    collect_specs: OnceCell<crate::Specs>,
202    resolve_crate: OnceCell<crate::ResolverOutput>,
203    desugar: Cache<LocalDefId, QueryResult<fhir::Node<'genv>>>,
204    fhir_crate: OnceCell<fhir::FluxItems<'genv>>,
205    lower_generics_of: Cache<DefId, ty::Generics<'tcx>>,
206    lower_predicates_of: Cache<DefId, QueryResult<ty::GenericPredicates>>,
207    lower_type_of: Cache<DefId, QueryResult<ty::EarlyBinder<ty::Ty>>>,
208    lower_fn_sig: Cache<DefId, QueryResult<ty::EarlyBinder<ty::PolyFnSig>>>,
209    normalized_defns: Cache<CrateNum, Rc<rty::NormalizedDefns>>,
210    func_sort: Cache<FluxDefId, rty::PolyFuncSort>,
211    qualifiers: OnceCell<QueryResult<Vec<rty::Qualifier>>>,
212    adt_sort_def_of: Cache<DefId, QueryResult<rty::AdtSortDef>>,
213    check_wf: Cache<LocalDefId, QueryResult<Rc<rty::WfckResults>>>,
214    adt_def: Cache<DefId, QueryResult<rty::AdtDef>>,
215    constant_info: Cache<DefId, QueryResult<rty::ConstantInfo>>,
216    generics_of: Cache<DefId, QueryResult<rty::Generics>>,
217    refinement_generics_of: Cache<DefId, QueryResult<rty::EarlyBinder<rty::RefinementGenerics>>>,
218    predicates_of: Cache<DefId, QueryResult<rty::EarlyBinder<rty::GenericPredicates>>>,
219    assoc_refinements_of: Cache<DefId, QueryResult<rty::AssocRefinements>>,
220    assoc_refinement_body: Cache<FluxDefId, QueryResult<rty::EarlyBinder<rty::Lambda>>>,
221    default_assoc_refinement_body:
222        Cache<FluxDefId, QueryResult<Option<rty::EarlyBinder<rty::Lambda>>>>,
223    sort_of_assoc_reft: Cache<FluxDefId, QueryResult<rty::EarlyBinder<rty::FuncSort>>>,
224    item_bounds: Cache<DefId, QueryResult<rty::EarlyBinder<List<rty::Clause>>>>,
225    type_of: Cache<DefId, QueryResult<rty::EarlyBinder<rty::TyOrCtor>>>,
226    variants_of: Cache<DefId, QueryResult<rty::Opaqueness<rty::EarlyBinder<rty::PolyVariants>>>>,
227    fn_sig: Cache<DefId, QueryResult<rty::EarlyBinder<rty::PolyFnSig>>>,
228    lower_late_bound_vars: Cache<LocalDefId, QueryResult<List<ty::BoundVariableKind>>>,
229}
230
231impl<'genv, 'tcx> Queries<'genv, 'tcx> {
232    pub(crate) fn new(providers: Providers) -> Self {
233        Self {
234            providers,
235            mir: Default::default(),
236            collect_specs: Default::default(),
237            resolve_crate: Default::default(),
238            desugar: Default::default(),
239            fhir_crate: Default::default(),
240            lower_generics_of: Default::default(),
241            lower_predicates_of: Default::default(),
242            lower_type_of: Default::default(),
243            lower_fn_sig: Default::default(),
244            normalized_defns: Default::default(),
245            func_sort: Default::default(),
246            qualifiers: Default::default(),
247            adt_sort_def_of: Default::default(),
248            check_wf: Default::default(),
249            adt_def: Default::default(),
250            constant_info: Default::default(),
251            generics_of: Default::default(),
252            refinement_generics_of: Default::default(),
253            predicates_of: Default::default(),
254            assoc_refinements_of: Default::default(),
255            assoc_refinement_body: Default::default(),
256            default_assoc_refinement_body: Default::default(),
257            sort_of_assoc_reft: Default::default(),
258            item_bounds: Default::default(),
259            type_of: Default::default(),
260            variants_of: Default::default(),
261            fn_sig: Default::default(),
262            lower_late_bound_vars: Default::default(),
263        }
264    }
265
266    pub(crate) fn mir(
267        &self,
268        genv: GlobalEnv<'genv, 'tcx>,
269        def_id: LocalDefId,
270    ) -> QueryResult<Rc<mir::Body<'tcx>>> {
271        run_with_cache(&self.mir, def_id, || {
272            let mir = unsafe { flux_common::mir_storage::retrieve_mir_body(genv.tcx(), def_id) };
273            let mir =
274                lowering::MirLoweringCtxt::lower_mir_body(genv.tcx(), genv.sess(), def_id, mir)?;
275            Ok(Rc::new(mir))
276        })
277    }
278
279    pub(crate) fn collect_specs(&'genv self, genv: GlobalEnv<'genv, 'tcx>) -> &'genv crate::Specs {
280        self.collect_specs
281            .get_or_init(|| (self.providers.collect_specs)(genv))
282    }
283
284    pub(crate) fn resolve_crate(
285        &'genv self,
286        genv: GlobalEnv<'genv, 'tcx>,
287    ) -> &'genv crate::ResolverOutput {
288        self.resolve_crate
289            .get_or_init(|| (self.providers.resolve_crate)(genv))
290    }
291
292    pub(crate) fn desugar(
293        &'genv self,
294        genv: GlobalEnv<'genv, 'tcx>,
295        def_id: LocalDefId,
296    ) -> QueryResult<fhir::Node<'genv>> {
297        if let Some(v) = self.desugar.borrow().get(&def_id) {
298            return v.clone();
299        }
300        match (self.providers.desugar)(genv, def_id) {
301            Ok(nodes) => {
302                let mut cache = self.desugar.borrow_mut();
303                cache.extend_unord(nodes.into_items().map(|(def_id, node)| (def_id, Ok(node))));
304                cache[&def_id].clone()
305            }
306            Err(err) => {
307                self.desugar.borrow_mut().insert(def_id, Err(err.clone()));
308                Err(err)
309            }
310        }
311    }
312
313    pub(crate) fn fhir_crate(
314        &'genv self,
315        genv: GlobalEnv<'genv, 'tcx>,
316    ) -> &'genv fhir::FluxItems<'genv> {
317        self.fhir_crate
318            .get_or_init(|| (self.providers.fhir_crate)(genv))
319    }
320
321    pub(crate) fn lower_generics_of(
322        &self,
323        genv: GlobalEnv<'genv, 'tcx>,
324        def_id: DefId,
325    ) -> ty::Generics<'tcx> {
326        run_with_cache(&self.lower_generics_of, def_id, || {
327            genv.tcx().generics_of(def_id).lower(genv.tcx())
328        })
329    }
330
331    pub(crate) fn lower_predicates_of(
332        &self,
333        genv: GlobalEnv,
334        def_id: DefId,
335    ) -> QueryResult<ty::GenericPredicates> {
336        run_with_cache(&self.lower_predicates_of, def_id, || {
337            genv.tcx()
338                .predicates_of(def_id)
339                .lower(genv.tcx())
340                .map_err(|err| QueryErr::unsupported(def_id, err))
341        })
342    }
343
344    pub(crate) fn lower_type_of(
345        &self,
346        genv: GlobalEnv,
347        def_id: DefId,
348    ) -> QueryResult<ty::EarlyBinder<ty::Ty>> {
349        run_with_cache(&self.lower_type_of, def_id, || {
350            let ty = genv.tcx().type_of(def_id).instantiate_identity();
351            Ok(ty::EarlyBinder(
352                ty.lower(genv.tcx())
353                    .map_err(|err| QueryErr::unsupported(def_id, err.into_err()))?,
354            ))
355        })
356    }
357
358    pub(crate) fn lower_fn_sig(
359        &self,
360        genv: GlobalEnv,
361        def_id: DefId,
362    ) -> QueryResult<ty::EarlyBinder<ty::PolyFnSig>> {
363        run_with_cache(&self.lower_fn_sig, def_id, || {
364            let fn_sig = genv.tcx().fn_sig(def_id).instantiate_identity();
365            Ok(ty::EarlyBinder(
366                fn_sig
367                    .lower(genv.tcx())
368                    .map_err(|err| QueryErr::unsupported(def_id, err.into_err()))?,
369            ))
370        })
371    }
372
373    pub(crate) fn lower_late_bound_vars(
374        &self,
375        genv: GlobalEnv,
376        def_id: LocalDefId,
377    ) -> QueryResult<List<ty::BoundVariableKind>> {
378        run_with_cache(&self.lower_late_bound_vars, def_id, || {
379            let hir_id = genv.tcx().local_def_id_to_hir_id(def_id);
380            genv.tcx()
381                .late_bound_vars(hir_id)
382                .lower(genv.tcx())
383                .map_err(|err| QueryErr::unsupported(def_id.to_def_id(), err.into_err()))
384        })
385    }
386
387    pub(crate) fn normalized_defns(
388        &self,
389        genv: GlobalEnv,
390        krate: CrateNum,
391    ) -> Rc<rty::NormalizedDefns> {
392        run_with_cache(&self.normalized_defns, krate, || {
393            if krate == LOCAL_CRATE {
394                Rc::new((self.providers.normalized_defns)(genv))
395            } else {
396                genv.cstore().normalized_defns(krate)
397            }
398        })
399    }
400
401    pub(crate) fn func_sort(&self, genv: GlobalEnv, def_id: FluxDefId) -> rty::PolyFuncSort {
402        run_with_cache(&self.func_sort, def_id, || {
403            dispatch_query_flux_id(
404                genv,
405                def_id,
406                |def_id| {
407                    // refinement functions cannot be extern specs so we simply grab the local id
408                    (self.providers.func_sort)(genv, def_id.local_id())
409                },
410                |def_id| genv.cstore().func_sort(def_id),
411                |_| {
412                    bug!(
413                        "cannot generate default function sort, the refinement must be defined somewhere"
414                    )
415                },
416            )
417        })
418    }
419
420    pub(crate) fn qualifiers(&self, genv: GlobalEnv) -> QueryResult<&[rty::Qualifier]> {
421        self.qualifiers
422            .get_or_init(|| (self.providers.qualifiers)(genv))
423            .as_deref()
424            .map_err(Clone::clone)
425    }
426
427    pub(crate) fn adt_sort_def_of(
428        &self,
429        genv: GlobalEnv,
430        def_id: DefId,
431    ) -> QueryResult<rty::AdtSortDef> {
432        run_with_cache(&self.adt_sort_def_of, def_id, || {
433            dispatch_query(
434                genv,
435                def_id,
436                |def_id| (self.providers.adt_sort_def_of)(genv, def_id.local_id()),
437                |def_id| genv.cstore().adt_sort_def(def_id),
438                |def_id| {
439                    let variants = IndexVec::from([rty::AdtSortVariant::new(vec![])]);
440                    Ok(rty::AdtSortDef::new(def_id, vec![], variants, false, true))
441                },
442            )
443        })
444    }
445
446    pub(crate) fn check_wf(
447        &self,
448        genv: GlobalEnv<'genv, '_>,
449        def_id: LocalDefId,
450    ) -> QueryResult<Rc<rty::WfckResults>> {
451        run_with_cache(&self.check_wf, def_id, || (self.providers.check_wf)(genv, def_id))
452    }
453
454    pub(crate) fn constant_info(
455        &self,
456        genv: GlobalEnv,
457        def_id: DefId,
458    ) -> QueryResult<rty::ConstantInfo> {
459        run_with_cache(&self.constant_info, def_id, || {
460            dispatch_query(
461                genv,
462                def_id,
463                |def_id| (self.providers.constant_info)(genv, def_id.local_id()),
464                |def_id| genv.cstore().constant_info(def_id),
465                |def_id| {
466                    // TODO(RJ): fix duplication with [`conv_constant`]` in `flux-fhir-analysis`
467                    let ty = genv.tcx().type_of(def_id).no_bound_vars().unwrap();
468                    if ty.is_integral() {
469                        let val = genv.tcx().const_eval_poly(def_id).ok().and_then(|val| {
470                            let val = val.try_to_scalar_int()?;
471                            rty::Constant::from_scalar_int(genv.tcx(), val, &ty)
472                        });
473                        if let Some(constant_) = val {
474                            return Ok(rty::ConstantInfo::Interpreted(
475                                rty::Expr::constant(constant_),
476                                rty::Sort::Int,
477                            ));
478                        }
479                    }
480                    Ok(rty::ConstantInfo::Uninterpreted)
481                },
482            )
483        })
484    }
485
486    pub(crate) fn adt_def(&self, genv: GlobalEnv, def_id: DefId) -> QueryResult<rty::AdtDef> {
487        run_with_cache(&self.adt_def, def_id, || {
488            dispatch_query(
489                genv,
490                def_id,
491                |def_id| (self.providers.adt_def)(genv, def_id.local_id()),
492                |def_id| genv.cstore().adt_def(def_id),
493                |def_id| {
494                    let adt_def = genv.tcx().adt_def(def_id).lower(genv.tcx());
495                    Ok(rty::AdtDef::new(adt_def, genv.adt_sort_def_of(def_id)?, vec![], false))
496                },
497            )
498        })
499    }
500
501    pub(crate) fn generics_of(&self, genv: GlobalEnv, def_id: DefId) -> QueryResult<rty::Generics> {
502        run_with_cache(&self.generics_of, def_id, || {
503            dispatch_query(
504                genv,
505                def_id,
506                |def_id| (self.providers.generics_of)(genv, def_id.local_id()),
507                |def_id| genv.cstore().generics_of(def_id),
508                |def_id| {
509                    Ok(refining::refine_generics(genv, def_id, &genv.lower_generics_of(def_id)))
510                },
511            )
512        })
513    }
514
515    pub(crate) fn refinement_generics_of(
516        &self,
517        genv: GlobalEnv,
518        def_id: DefId,
519    ) -> QueryResult<rty::EarlyBinder<rty::RefinementGenerics>> {
520        run_with_cache(&self.refinement_generics_of, def_id, || {
521            dispatch_query(
522                genv,
523                def_id,
524                |def_id| (self.providers.refinement_generics_of)(genv, def_id.local_id()),
525                |def_id| genv.cstore().refinement_generics_of(def_id),
526                |def_id| {
527                    let parent = genv.tcx().generics_of(def_id).parent;
528                    Ok(rty::EarlyBinder(rty::RefinementGenerics {
529                        parent,
530                        parent_count: 0,
531                        own_params: List::empty(),
532                    }))
533                },
534            )
535        })
536    }
537
538    pub(crate) fn item_bounds(
539        &self,
540        genv: GlobalEnv<'genv, 'tcx>,
541        def_id: DefId,
542    ) -> QueryResult<rty::EarlyBinder<List<rty::Clause>>> {
543        run_with_cache(&self.item_bounds, def_id, || {
544            dispatch_query(
545                genv,
546                def_id,
547                |def_id| (self.providers.item_bounds)(genv, def_id.local_id()),
548                |def_id| genv.cstore().item_bounds(def_id),
549                |def_id| {
550                    let clauses = genv
551                        .tcx()
552                        .item_bounds(def_id)
553                        .skip_binder()
554                        .lower(genv.tcx())
555                        .map_err(|err| QueryErr::unsupported(def_id, err))?
556                        .refine(&Refiner::default_for_item(genv, def_id)?)?;
557
558                    Ok(rty::EarlyBinder(clauses))
559                },
560            )
561        })
562    }
563
564    pub(crate) fn predicates_of(
565        &self,
566        genv: GlobalEnv,
567        def_id: DefId,
568    ) -> QueryResult<rty::EarlyBinder<rty::GenericPredicates>> {
569        run_with_cache(&self.predicates_of, def_id, || {
570            dispatch_query(
571                genv,
572                def_id,
573                |def_id| (self.providers.predicates_of)(genv, def_id.local_id()),
574                |def_id| genv.cstore().predicates_of(def_id),
575                |def_id| {
576                    let predicates = genv
577                        .lower_predicates_of(def_id)?
578                        .refine(&Refiner::default_for_item(genv, def_id)?)?;
579                    Ok(rty::EarlyBinder(predicates))
580                },
581            )
582        })
583    }
584
585    pub(crate) fn assoc_refinements_of(
586        &self,
587        genv: GlobalEnv,
588        def_id: DefId,
589    ) -> QueryResult<rty::AssocRefinements> {
590        run_with_cache(&self.assoc_refinements_of, def_id, || {
591            dispatch_query(
592                genv,
593                def_id,
594                |def_id| (self.providers.assoc_refinements_of)(genv, def_id.local_id()),
595                |def_id| genv.cstore().assoc_refinements_of(def_id),
596                |_| Ok(rty::AssocRefinements::default()),
597            )
598        })
599    }
600
601    pub(crate) fn assoc_refinement_body(
602        &self,
603        genv: GlobalEnv,
604        impl_assoc_id: FluxDefId,
605    ) -> QueryResult<rty::EarlyBinder<rty::Lambda>> {
606        run_with_cache(&self.assoc_refinement_body, impl_assoc_id, || {
607            dispatch_query_flux_id(
608                genv,
609                impl_assoc_id,
610                |impl_assoc_id| (self.providers.assoc_refinement_body)(genv, impl_assoc_id),
611                |impl_assoc_id| genv.cstore().assoc_refinements_def(impl_assoc_id),
612                |impl_assoc_id| {
613                    Err(query_bug!(
614                        impl_assoc_id.parent(),
615                        "cannot generate default associate refinement for extern impl"
616                    ))
617                },
618            )
619        })
620    }
621
622    pub(crate) fn default_assoc_refinement_body(
623        &self,
624        genv: GlobalEnv,
625        trait_assoc_id: FluxDefId,
626    ) -> QueryResult<Option<rty::EarlyBinder<rty::Lambda>>> {
627        run_with_cache(&self.default_assoc_refinement_body, trait_assoc_id, || {
628            dispatch_query_flux_id(
629                genv,
630                trait_assoc_id,
631                |trait_assoc_id| {
632                    (self.providers.default_assoc_refinement_body)(genv, trait_assoc_id)
633                },
634                |trait_assoc_id| genv.cstore().default_assoc_refinements_def(trait_assoc_id),
635                |trait_assoc_id| {
636                    Err(query_bug!(
637                        trait_assoc_id.parent(),
638                        "cannot generate default assoc refinement for extern trait"
639                    ))
640                },
641            )
642        })
643    }
644
645    pub(crate) fn sort_of_assoc_reft(
646        &self,
647        genv: GlobalEnv,
648        assoc_id: FluxDefId,
649    ) -> QueryResult<rty::EarlyBinder<rty::FuncSort>> {
650        run_with_cache(&self.sort_of_assoc_reft, assoc_id, || {
651            dispatch_query_flux_id(
652                genv,
653                assoc_id,
654                |assoc_id| (self.providers.sort_of_assoc_reft)(genv, assoc_id),
655                |assoc_id| genv.cstore().sort_of_assoc_reft(assoc_id),
656                |assoc_id| {
657                    Err(query_bug!(
658                        assoc_id.parent(),
659                        "cannot generate default sort for assoc refinement in extern crate"
660                    ))
661                },
662            )
663        })
664    }
665
666    pub(crate) fn type_of(
667        &self,
668        genv: GlobalEnv,
669        def_id: DefId,
670    ) -> QueryResult<rty::EarlyBinder<rty::TyOrCtor>> {
671        run_with_cache(&self.type_of, def_id, || {
672            dispatch_query(
673                genv,
674                def_id,
675                |def_id| (self.providers.type_of)(genv, def_id.local_id()),
676                |def_id| genv.cstore().type_of(def_id),
677                |def_id| {
678                    // If we're given a type parameter, provide the generics of the parent container.
679                    let generics_def_id = match genv.def_kind(def_id) {
680                        DefKind::TyParam => genv.tcx().parent(def_id),
681                        _ => def_id,
682                    };
683                    let ty = genv.lower_type_of(def_id)?.skip_binder();
684                    Ok(rty::EarlyBinder(
685                        Refiner::default_for_item(genv, generics_def_id)?
686                            .refine_ty_or_base(&ty)?
687                            .into(),
688                    ))
689                },
690            )
691        })
692    }
693
694    pub(crate) fn variants_of(
695        &self,
696        genv: GlobalEnv,
697        def_id: DefId,
698    ) -> QueryResult<rty::Opaqueness<rty::EarlyBinder<rty::PolyVariants>>> {
699        run_with_cache(&self.variants_of, def_id, || {
700            dispatch_query(
701                genv,
702                def_id,
703                |def_id| (self.providers.variants_of)(genv, def_id.local_id()),
704                |def_id| genv.cstore().variants(def_id),
705                |def_id| {
706                    let variants = genv
707                        .tcx()
708                        .adt_def(def_id)
709                        .variants()
710                        .indices()
711                        .map(|variant_idx| {
712                            Refiner::default_for_item(genv, def_id)?
713                                .refine_variant_def(def_id, variant_idx)
714                        })
715                        .try_collect()?;
716                    Ok(rty::Opaqueness::Transparent(rty::EarlyBinder(variants)))
717                },
718            )
719        })
720    }
721
722    pub(crate) fn fn_sig(
723        &self,
724        genv: GlobalEnv,
725        def_id: DefId,
726    ) -> QueryResult<rty::EarlyBinder<rty::PolyFnSig>> {
727        run_with_cache(&self.fn_sig, def_id, || {
728            dispatch_query(
729                genv,
730                def_id,
731                |def_id| (self.providers.fn_sig)(genv, def_id.local_id()),
732                |def_id| genv.cstore().fn_sig(def_id),
733                |def_id| {
734                    let fn_sig = genv
735                        .lower_fn_sig(def_id)?
736                        .skip_binder()
737                        .refine(&Refiner::default_for_item(genv, def_id)?)?
738                        .hoist_input_binders();
739                    Ok(rty::EarlyBinder(fn_sig))
740                },
741            )
742        })
743    }
744}
745
746/// [Resolve] the `def_id` and *dispatch* it to a provider (`local`, `external`, or `default`).
747///
748/// [Resolve]: GlobalEnv::resolve_id
749fn dispatch_query<R>(
750    genv: GlobalEnv,
751    def_id: DefId,
752    local: impl FnOnce(MaybeExternId) -> R,
753    external: impl FnOnce(DefId) -> Option<R>,
754    default: impl FnOnce(DefId) -> R,
755) -> R {
756    match genv.resolve_id(def_id) {
757        ResolvedDefId::Local(local_id) => {
758            // Case 1: `def_id` is a `LocalDefId` so forward it to the *local provider*
759            local(MaybeExternId::Local(local_id))
760        }
761        ResolvedDefId::ExternSpec(local_id, def_id) => {
762            // Case 2: `def_id` is a `LocalDefId` wrapping an extern spec, so we also
763            // forward it to the local provider
764            local(MaybeExternId::Extern(local_id, def_id))
765        }
766        ResolvedDefId::Extern(def_id) if let Some(v) = external(def_id) => {
767            // Case 3: `def_id` is an external `def_id` for which we have an annotation in the
768            // *external provider*
769            v
770        }
771        ResolvedDefId::Extern(def_id) => {
772            // Case 4: If none of the above, we generate a default annotation
773            default(def_id)
774        }
775    }
776}
777
778fn dispatch_query_flux_id<R>(
779    genv: GlobalEnv,
780    def_id: FluxId<DefId>,
781    local: impl FnOnce(FluxId<MaybeExternId>) -> R,
782    external: impl FnOnce(FluxId<DefId>) -> Option<R>,
783    default: impl FnOnce(FluxId<DefId>) -> R,
784) -> R {
785    #[allow(
786        clippy::disallowed_methods,
787        reason = "we are mapping the parent id to a different representation which still guarantees the existence of the item"
788    )]
789    dispatch_query(
790        genv,
791        def_id.parent(),
792        |container_id| local(FluxId::new(container_id, def_id.name())),
793        |container_id| external(FluxId::new(container_id, def_id.name())),
794        |container_id| default(FluxId::new(container_id, def_id.name())),
795    )
796}
797
798fn run_with_cache<K, V>(cache: &Cache<K, V>, key: K, f: impl FnOnce() -> V) -> V
799where
800    K: std::hash::Hash + Eq,
801    V: Clone,
802{
803    if let Some(v) = cache.borrow().get(&key) {
804        return v.clone();
805    }
806    let v = f();
807    cache.borrow_mut().insert(key, v.clone());
808    v
809}
810
811impl<'a> Diagnostic<'a> for QueryErr {
812    #[track_caller]
813    fn into_diag(
814        self,
815        dcx: rustc_errors::DiagCtxtHandle<'a>,
816        _level: rustc_errors::Level,
817    ) -> rustc_errors::Diag<'a, ErrorGuaranteed> {
818        use crate::fluent_generated as fluent;
819
820        rustc_middle::ty::tls::with_opt(
821            #[track_caller]
822            |tcx| {
823                let tcx = tcx.expect("no TyCtxt stored in tls");
824                match self {
825                    QueryErr::Unsupported { def_id, err } => {
826                        let span = err.span.unwrap_or_else(|| tcx.def_span(def_id));
827                        let mut diag = dcx.struct_span_err(span, fluent::middle_query_unsupported);
828                        diag.code(E0999);
829                        diag.note(err.descr);
830                        diag
831                    }
832                    QueryErr::Ignored { def_id } => {
833                        let def_span = tcx.def_span(def_id);
834                        let mut diag =
835                            dcx.struct_span_err(def_span, fluent::middle_query_ignored_item);
836                        diag.code(E0999);
837                        diag
838                    }
839                    QueryErr::InvalidGenericArg { def_id } => {
840                        let def_span = tcx.def_span(def_id);
841                        let mut diag =
842                            dcx.struct_span_err(def_span, fluent::middle_query_invalid_generic_arg);
843                        diag.code(E0999);
844                        diag
845                    }
846                    QueryErr::MissingAssocReft { impl_id, name, .. } => {
847                        let def_span = tcx.def_span(impl_id);
848                        let mut diag =
849                            dcx.struct_span_err(def_span, fluent::middle_query_missing_assoc_reft);
850                        diag.arg("name", name);
851                        diag.code(E0999);
852                        diag
853                    }
854                    QueryErr::Bug { def_id, location, msg } => {
855                        let mut diag = dcx.struct_err(fluent::middle_query_bug);
856                        if let Some(def_id) = def_id {
857                            diag.span(tcx.def_span(def_id));
858                        }
859                        diag.arg("location", location);
860                        diag.note(msg);
861                        diag
862                    }
863                    QueryErr::Emitted(_) => {
864                        let mut diag = dcx.struct_err("QueryErr::Emitted should be emitted");
865                        diag.downgrade_to_delayed_bug();
866                        diag
867                    }
868                }
869            },
870        )
871    }
872}
873
874impl<'a> Diagnostic<'a> for QueryErrAt {
875    #[track_caller]
876    fn into_diag(
877        self,
878        dcx: rustc_errors::DiagCtxtHandle<'a>,
879        level: rustc_errors::Level,
880    ) -> rustc_errors::Diag<'a, ErrorGuaranteed> {
881        use crate::fluent_generated as fluent;
882
883        rustc_middle::ty::tls::with_opt(
884            #[track_caller]
885            |tcx| {
886                let tcx = tcx.expect("no TyCtxt stored in tls");
887                let mut diag = match self.err {
888                    QueryErr::Unsupported { def_id, err, .. } => {
889                        let mut diag =
890                            dcx.struct_span_err(self.span, fluent::middle_query_unsupported_at);
891                        diag.arg("kind", tcx.def_kind(def_id).descr(def_id));
892                        if let Some(def_ident_span) = tcx.def_ident_span(def_id) {
893                            diag.span_note(def_ident_span, fluent::_subdiag::note);
894                        }
895                        diag.note(err.descr);
896                        diag
897                    }
898                    QueryErr::Ignored { def_id } => {
899                        let mut diag =
900                            dcx.struct_span_err(self.span, fluent::middle_query_ignored_at);
901                        diag.arg("kind", tcx.def_kind(def_id).descr(def_id));
902                        diag.arg("name", def_id_to_string(def_id));
903                        diag.span_label(self.span, fluent::_subdiag::label);
904                        diag
905                    }
906                    QueryErr::MissingAssocReft { name, .. } => {
907                        let mut diag = dcx
908                            .struct_span_err(self.span, fluent::middle_query_missing_assoc_reft_at);
909                        diag.arg("name", name);
910                        diag.code(E0999);
911                        diag
912                    }
913                    QueryErr::InvalidGenericArg { .. }
914                    | QueryErr::Emitted(_)
915                    | QueryErr::Bug { .. } => {
916                        let mut diag = self.err.into_diag(dcx, level);
917                        diag.span(self.span);
918                        diag
919                    }
920                };
921                diag.code(E0999);
922                diag
923            },
924        )
925    }
926}
927
928impl From<ErrorGuaranteed> for QueryErr {
929    fn from(err: ErrorGuaranteed) -> Self {
930        Self::Emitted(err)
931    }
932}
933
934pub fn try_query<T>(f: impl FnOnce() -> QueryResult<T>) -> QueryResult<T> {
935    f()
936}