1use std::{alloc, ptr, rc::Rc, slice};
2
3use flux_arc_interner::List;
4use flux_common::{bug, result::ErrorEmitter};
5use flux_config as config;
6use flux_errors::FluxSession;
7use flux_rustc_bridge::{self, lowering::Lower, mir, ty};
8use rustc_data_structures::unord::UnordSet;
9use rustc_hir::{
10    def::DefKind,
11    def_id::{CrateNum, DefId, LocalDefId},
12};
13use rustc_middle::{
14    query::IntoQueryParam,
15    ty::{TyCtxt, Variance},
16};
17use rustc_span::Span;
18pub use rustc_span::{Symbol, symbol::Ident};
19
20use crate::{
21    cstore::CrateStoreDyn,
22    def_id::{FluxDefId, FluxLocalDefId, MaybeExternId, ResolvedDefId},
23    fhir::{self, VariantIdx},
24    queries::{Providers, Queries, QueryErr, QueryResult},
25    query_bug,
26    rty::{
27        self,
28        refining::{Refine as _, Refiner},
29    },
30};
31
32#[derive(Clone, Copy)]
33pub struct GlobalEnv<'genv, 'tcx> {
34    inner: &'genv GlobalEnvInner<'genv, 'tcx>,
35}
36
37struct GlobalEnvInner<'genv, 'tcx> {
38    tcx: TyCtxt<'tcx>,
39    sess: &'genv FluxSession,
40    arena: &'genv fhir::Arena,
41    cstore: Box<CrateStoreDyn>,
42    queries: Queries<'genv, 'tcx>,
43}
44
45impl<'tcx> GlobalEnv<'_, 'tcx> {
46    pub fn enter<'a, R>(
47        tcx: TyCtxt<'tcx>,
48        sess: &'a FluxSession,
49        cstore: Box<CrateStoreDyn>,
50        arena: &'a fhir::Arena,
51        providers: Providers,
52        f: impl for<'genv> FnOnce(GlobalEnv<'genv, 'tcx>) -> R,
53    ) -> R {
54        let inner = GlobalEnvInner { tcx, sess, cstore, arena, queries: Queries::new(providers) };
55        f(GlobalEnv { inner: &inner })
56    }
57}
58
59impl<'genv, 'tcx> GlobalEnv<'genv, 'tcx> {
60    pub fn tcx(self) -> TyCtxt<'tcx> {
61        self.inner.tcx
62    }
63
64    pub fn sess(self) -> &'genv FluxSession {
65        self.inner.sess
66    }
67
68    pub fn collect_specs(self) -> &'genv crate::Specs {
69        self.inner.queries.collect_specs(self)
70    }
71
72    pub fn resolve_crate(self) -> &'genv crate::ResolverOutput {
73        self.inner.queries.resolve_crate(self)
74    }
75
76    pub fn desugar(self, def_id: LocalDefId) -> QueryResult<fhir::Node<'genv>> {
77        self.inner.queries.desugar(self, def_id)
78    }
79
80    pub fn fhir_attr_map(self, def_id: LocalDefId) -> fhir::AttrMap<'genv> {
81        self.inner.queries.fhir_attr_map(self, def_id)
82    }
83
84    pub fn fhir_crate(self) -> &'genv fhir::FluxItems<'genv> {
85        self.inner.queries.fhir_crate(self)
86    }
87
88    pub fn alloc<T>(&self, val: T) -> &'genv T {
89        self.inner.arena.alloc(val)
90    }
91
92    pub fn alloc_slice<T: Copy>(self, slice: &[T]) -> &'genv [T] {
93        self.inner.arena.alloc_slice_copy(slice)
94    }
95
96    pub fn alloc_slice_fill_iter<T, I>(self, it: I) -> &'genv [T]
97    where
98        I: IntoIterator<Item = T>,
99        I::IntoIter: ExactSizeIterator,
100    {
101        self.inner.arena.alloc_slice_fill_iter(it)
102    }
103
104    pub fn def_kind(&self, def_id: impl IntoQueryParam<DefId>) -> DefKind {
105        self.tcx().def_kind(def_id.into_query_param())
106    }
107
108    pub fn alloc_slice_with_capacity<T, I>(self, cap: usize, it: I) -> &'genv [T]
118    where
119        I: IntoIterator<Item = T>,
120    {
121        let layout = alloc::Layout::array::<T>(cap).unwrap_or_else(|_| bug!("out of memory"));
122        let dst = self.inner.arena.alloc_layout(layout).cast::<T>();
123        unsafe {
124            let mut len = 0;
125            for (i, v) in it.into_iter().take(cap).enumerate() {
126                len += 1;
127                ptr::write(dst.as_ptr().add(i), v);
128            }
129
130            slice::from_raw_parts(dst.as_ptr(), len)
131        }
132    }
133
134    pub fn normalized_info(self, did: FluxDefId) -> rty::NormalizeInfo {
135        self.normalized_defns(did.krate()).func_info(did).clone()
136    }
137
138    pub fn normalized_defns(self, krate: CrateNum) -> Rc<rty::NormalizedDefns> {
139        self.inner.queries.normalized_defns(self, krate)
140    }
141
142    pub fn prim_rel_for(self, op: &rty::BinOp) -> QueryResult<Option<&'genv rty::PrimRel>> {
143        Ok(self.inner.queries.prim_rel(self)?.get(op))
144    }
145
146    pub fn qualifiers(self) -> QueryResult<&'genv [rty::Qualifier]> {
147        self.inner.queries.qualifiers(self)
148    }
149
150    pub fn qualifiers_for(
152        self,
153        did: LocalDefId,
154    ) -> QueryResult<impl Iterator<Item = &'genv rty::Qualifier>> {
155        let quals = self.fhir_attr_map(did).qualifiers;
156        let names: UnordSet<_> = quals.iter().copied().collect();
157        Ok(self
158            .qualifiers()?
159            .iter()
160            .filter(move |qual| qual.global || names.contains(&qual.def_id)))
161    }
162
163    pub fn reveals_for(self, did: LocalDefId) -> &'genv [FluxDefId] {
165        self.fhir_attr_map(did).reveals
166    }
167
168    pub fn func_sort(self, def_id: impl IntoQueryParam<FluxDefId>) -> rty::PolyFuncSort {
169        self.inner
170            .queries
171            .func_sort(self, def_id.into_query_param())
172    }
173
174    pub fn func_span(self, def_id: impl IntoQueryParam<FluxDefId>) -> Span {
175        self.inner
176            .queries
177            .func_span(self, def_id.into_query_param())
178    }
179
180    pub fn should_inline_fun(self, def_id: FluxDefId) -> bool {
181        let is_poly = self.func_sort(def_id).params().len() > 0;
182        is_poly || !flux_config::smt_define_fun()
183    }
184
185    pub fn variances_of(self, did: DefId) -> &'tcx [Variance] {
186        self.tcx().variances_of(did)
187    }
188
189    pub fn mir(self, def_id: LocalDefId) -> QueryResult<Rc<mir::Body<'tcx>>> {
190        self.inner.queries.mir(self, def_id)
191    }
192
193    pub fn lower_generics_of(self, def_id: impl IntoQueryParam<DefId>) -> ty::Generics<'tcx> {
194        self.inner
195            .queries
196            .lower_generics_of(self, def_id.into_query_param())
197    }
198
199    pub fn lower_predicates_of(
200        self,
201        def_id: impl IntoQueryParam<DefId>,
202    ) -> QueryResult<ty::GenericPredicates> {
203        self.inner
204            .queries
205            .lower_predicates_of(self, def_id.into_query_param())
206    }
207
208    pub fn lower_type_of(
209        self,
210        def_id: impl IntoQueryParam<DefId>,
211    ) -> QueryResult<ty::EarlyBinder<ty::Ty>> {
212        self.inner
213            .queries
214            .lower_type_of(self, def_id.into_query_param())
215    }
216
217    pub fn lower_fn_sig(
218        self,
219        def_id: impl Into<DefId>,
220    ) -> QueryResult<ty::EarlyBinder<ty::PolyFnSig>> {
221        self.inner.queries.lower_fn_sig(self, def_id.into())
222    }
223
224    pub fn adt_def(self, def_id: impl IntoQueryParam<DefId>) -> QueryResult<rty::AdtDef> {
225        self.inner.queries.adt_def(self, def_id.into_query_param())
226    }
227
228    pub fn constant_info(
229        self,
230        def_id: impl IntoQueryParam<DefId>,
231    ) -> QueryResult<rty::ConstantInfo> {
232        self.inner
233            .queries
234            .constant_info(self, def_id.into_query_param())
235    }
236
237    pub fn adt_sort_def_of(
238        self,
239        def_id: impl IntoQueryParam<DefId>,
240    ) -> QueryResult<rty::AdtSortDef> {
241        self.inner
242            .queries
243            .adt_sort_def_of(self, def_id.into_query_param())
244    }
245
246    pub fn sort_decl_param_count(self, def_id: impl IntoQueryParam<FluxDefId>) -> usize {
247        self.inner
248            .queries
249            .sort_decl_param_count(self, def_id.into_query_param())
250    }
251
252    pub fn check_wf(self, def_id: LocalDefId) -> QueryResult<Rc<rty::WfckResults>> {
253        self.inner.queries.check_wf(self, def_id)
254    }
255
256    pub fn impl_trait_ref(
257        self,
258        impl_id: DefId,
259    ) -> QueryResult<Option<rty::EarlyBinder<rty::TraitRef>>> {
260        let Some(trait_ref) = self.tcx().impl_trait_ref(impl_id) else { return Ok(None) };
261        let trait_ref = trait_ref.skip_binder();
262        let trait_ref = trait_ref
263            .lower(self.tcx())
264            .map_err(|err| QueryErr::unsupported(trait_ref.def_id, err.into_err()))?
265            .refine(&Refiner::default_for_item(self, impl_id)?)?;
266        Ok(Some(rty::EarlyBinder(trait_ref)))
267    }
268
269    pub fn generics_of(self, def_id: impl IntoQueryParam<DefId>) -> QueryResult<rty::Generics> {
270        self.inner
271            .queries
272            .generics_of(self, def_id.into_query_param())
273    }
274
275    pub fn refinement_generics_of(
276        self,
277        def_id: impl IntoQueryParam<DefId>,
278    ) -> QueryResult<rty::EarlyBinder<rty::RefinementGenerics>> {
279        self.inner
280            .queries
281            .refinement_generics_of(self, def_id.into_query_param())
282    }
283
284    pub fn predicates_of(
285        self,
286        def_id: impl IntoQueryParam<DefId>,
287    ) -> QueryResult<rty::EarlyBinder<rty::GenericPredicates>> {
288        self.inner
289            .queries
290            .predicates_of(self, def_id.into_query_param())
291    }
292
293    pub fn assoc_refinements_of(
294        self,
295        def_id: impl IntoQueryParam<DefId>,
296    ) -> QueryResult<rty::AssocRefinements> {
297        self.inner
298            .queries
299            .assoc_refinements_of(self, def_id.into_query_param())
300    }
301
302    pub fn assoc_refinement(self, assoc_id: FluxDefId) -> QueryResult<rty::AssocReft> {
303        Ok(self.assoc_refinements_of(assoc_id.parent())?.get(assoc_id))
304    }
305
306    pub fn assoc_refinement_body_for_impl(
314        self,
315        trait_assoc_id: FluxDefId,
316        impl_id: DefId,
317    ) -> QueryResult<rty::EarlyBinder<rty::Lambda>> {
318        let impl_assoc_refts = self.assoc_refinements_of(impl_id)?;
320        if let Some(impl_assoc_reft) = impl_assoc_refts.find(trait_assoc_id.name()) {
321            return self.assoc_refinement_body(impl_assoc_reft.def_id());
322        }
323
324        if let Some(body) = self.default_assoc_refinement_body(trait_assoc_id)? {
326            let impl_trait_ref = self
327                .impl_trait_ref(impl_id)?
328                .unwrap()
329                .instantiate_identity();
330            return Ok(rty::EarlyBinder(body.instantiate(self.tcx(), &impl_trait_ref.args, &[])));
331        }
332
333        Err(QueryErr::MissingAssocReft {
334            impl_id,
335            trait_id: trait_assoc_id.parent(),
336            name: trait_assoc_id.name(),
337        })
338    }
339
340    pub fn default_assoc_refinement_body(
341        self,
342        trait_assoc_id: FluxDefId,
343    ) -> QueryResult<Option<rty::EarlyBinder<rty::Lambda>>> {
344        self.inner
345            .queries
346            .default_assoc_refinement_body(self, trait_assoc_id)
347    }
348
349    pub fn assoc_refinement_body(
350        self,
351        impl_assoc_id: FluxDefId,
352    ) -> QueryResult<rty::EarlyBinder<rty::Lambda>> {
353        self.inner
354            .queries
355            .assoc_refinement_body(self, impl_assoc_id)
356    }
357
358    pub fn sort_of_assoc_reft(
359        self,
360        assoc_id: FluxDefId,
361    ) -> QueryResult<rty::EarlyBinder<rty::FuncSort>> {
362        self.inner.queries.sort_of_assoc_reft(self, assoc_id)
363    }
364
365    pub fn item_bounds(self, def_id: DefId) -> QueryResult<rty::EarlyBinder<List<rty::Clause>>> {
366        self.inner.queries.item_bounds(self, def_id)
367    }
368
369    pub fn type_of(
370        self,
371        def_id: impl IntoQueryParam<DefId>,
372    ) -> QueryResult<rty::EarlyBinder<rty::TyOrCtor>> {
373        self.inner.queries.type_of(self, def_id.into_query_param())
374    }
375
376    pub fn fn_sig(
377        self,
378        def_id: impl IntoQueryParam<DefId>,
379    ) -> QueryResult<rty::EarlyBinder<rty::PolyFnSig>> {
380        self.inner.queries.fn_sig(self, def_id.into_query_param())
381    }
382
383    pub fn variants_of(
384        self,
385        def_id: impl IntoQueryParam<DefId>,
386    ) -> QueryResult<rty::Opaqueness<rty::EarlyBinder<rty::PolyVariants>>> {
387        self.inner
388            .queries
389            .variants_of(self, def_id.into_query_param())
390    }
391
392    pub fn variant_sig(
393        self,
394        def_id: DefId,
395        variant_idx: VariantIdx,
396    ) -> QueryResult<rty::Opaqueness<rty::EarlyBinder<rty::PolyVariant>>> {
397        Ok(self
398            .variants_of(def_id)?
399            .map(|variants| variants.map(|variants| variants[variant_idx.as_usize()].clone())))
400    }
401
402    pub fn lower_late_bound_vars(
403        self,
404        def_id: LocalDefId,
405    ) -> QueryResult<List<ty::BoundVariableKind>> {
406        self.inner.queries.lower_late_bound_vars(self, def_id)
407    }
408
409    pub fn is_box(&self, res: fhir::Res) -> bool {
410        res.is_box(self.tcx())
411    }
412
413    pub fn def_id_to_param_index(&self, def_id: DefId) -> u32 {
414        let parent = self.tcx().parent(def_id);
415        let generics = self.tcx().generics_of(parent);
416        generics.param_def_id_to_index(self.tcx(), def_id).unwrap()
417    }
418
419    pub(crate) fn cstore(self) -> &'genv CrateStoreDyn {
420        &*self.inner.cstore
421    }
422
423    pub fn has_trusted_impl(&self, def_id: DefId) -> bool {
424        if let Some(did) = self
425            .resolve_id(def_id)
426            .as_maybe_extern()
427            .map(|id| id.local_id())
428        {
429            self.trusted_impl(did)
430        } else {
431            false
432        }
433    }
434
435    pub fn is_fn_output(&self, def_id: DefId) -> bool {
439        let def_span = self.tcx().def_span(def_id);
440        self.tcx()
441            .require_lang_item(rustc_hir::LangItem::FnOnceOutput, def_span)
442            == def_id
443    }
444
445    pub fn iter_local_def_id(self) -> impl Iterator<Item = LocalDefId> + use<'tcx, 'genv> {
447        self.tcx().iter_local_def_id().filter(move |&local_def_id| {
448            self.maybe_extern_id(local_def_id).is_local() && !self.is_dummy(local_def_id)
449        })
450    }
451
452    pub fn iter_extern_def_id(self) -> impl Iterator<Item = DefId> + use<'tcx, 'genv> {
453        self.tcx()
454            .iter_local_def_id()
455            .filter_map(move |local_def_id| self.maybe_extern_id(local_def_id).as_extern())
456    }
457
458    pub fn maybe_extern_id(self, local_id: LocalDefId) -> MaybeExternId {
459        self.collect_specs()
460            .local_id_to_extern_id
461            .get(&local_id)
462            .map_or_else(
463                || MaybeExternId::Local(local_id),
464                |def_id| MaybeExternId::Extern(local_id, *def_id),
465            )
466    }
467
468    #[expect(clippy::disallowed_methods)]
469    pub fn resolve_id(self, def_id: DefId) -> ResolvedDefId {
470        let maybe_extern_spec = self
471            .collect_specs()
472            .extern_id_to_local_id
473            .get(&def_id)
474            .copied();
475        if let Some(local_id) = maybe_extern_spec {
476            ResolvedDefId::ExternSpec(local_id, def_id)
477        } else if let Some(local_id) = def_id.as_local() {
478            debug_assert!(
479                self.maybe_extern_id(local_id).is_local(),
480                "def id points to dummy local item `{def_id:?}`"
481            );
482            ResolvedDefId::Local(local_id)
483        } else {
484            ResolvedDefId::Extern(def_id)
485        }
486    }
487
488    pub fn infer_opts(self, def_id: LocalDefId) -> config::InferOpts {
489        let mut opts = config::PartialInferOpts::default();
490        self.traverse_parents(def_id, |did| {
491            if let Some(o) = self.fhir_attr_map(did).infer_opts() {
492                opts.merge(&o);
493            }
494            None::<!>
495        });
496        opts.into()
497    }
498
499    pub fn trusted(self, def_id: LocalDefId) -> bool {
503        self.traverse_parents(def_id, |did| self.fhir_attr_map(did).trusted())
504            .map(|trusted| trusted.to_bool())
505            .unwrap_or_else(config::trusted_default)
506    }
507
508    pub fn trusted_impl(self, def_id: LocalDefId) -> bool {
509        self.traverse_parents(def_id, |did| self.fhir_attr_map(did).trusted_impl())
510            .map(|trusted| trusted.to_bool())
511            .unwrap_or(false)
512    }
513
514    pub fn is_dummy(self, def_id: LocalDefId) -> bool {
518        self.traverse_parents(def_id, |did| {
519            self.collect_specs()
520                .dummy_extern
521                .contains(&did)
522                .then_some(())
523        })
524        .is_some()
525    }
526
527    pub fn ignored(self, def_id: LocalDefId) -> bool {
531        self.traverse_parents(def_id, |did| self.fhir_attr_map(did).ignored())
532            .map(|ignored| ignored.to_bool())
533            .unwrap_or_else(config::ignore_default)
534    }
535
536    pub fn should_fail(self, def_id: LocalDefId) -> bool {
538        self.fhir_attr_map(def_id).should_fail()
539    }
540
541    pub fn proven_externally(self, def_id: LocalDefId) -> bool {
543        self.fhir_attr_map(def_id).proven_externally()
544    }
545
546    fn traverse_parents<T>(
548        self,
549        mut def_id: LocalDefId,
550        mut f: impl FnMut(LocalDefId) -> Option<T>,
551    ) -> Option<T> {
552        loop {
553            if let Some(v) = f(def_id) {
554                break Some(v);
555            }
556
557            if let Some(parent) = self.tcx().opt_local_parent(def_id) {
558                def_id = parent;
559            } else {
560                break None;
561            }
562        }
563    }
564}
565
566impl<'genv, 'tcx> GlobalEnv<'genv, 'tcx> {
567    pub fn fhir_iter_flux_items(
568        self,
569    ) -> impl Iterator<Item = (FluxLocalDefId, fhir::FluxItem<'genv>)> {
570        self.fhir_crate()
571            .items
572            .iter()
573            .map(|(id, item)| (*id, *item))
574    }
575
576    pub fn fhir_sort_decl(&self, def_id: FluxLocalDefId) -> Option<&fhir::SortDecl> {
577        self.fhir_crate().items.get(&def_id).and_then(|item| {
578            if let fhir::FluxItem::SortDecl(sort_decl) = item { Some(*sort_decl) } else { None }
579        })
580    }
581
582    pub fn fhir_spec_func_body(
583        &self,
584        def_id: FluxLocalDefId,
585    ) -> Option<&'genv fhir::SpecFunc<'genv>> {
586        self.fhir_crate()
587            .items
588            .get(&def_id)
589            .and_then(|item| if let fhir::FluxItem::Func(defn) = item { Some(*defn) } else { None })
590    }
591
592    pub fn fhir_qualifiers(self) -> impl Iterator<Item = &'genv fhir::Qualifier<'genv>> {
593        self.fhir_crate().items.values().filter_map(|item| {
594            if let fhir::FluxItem::Qualifier(qual) = item { Some(*qual) } else { None }
595        })
596    }
597
598    pub fn fhir_primop_props(self) -> impl Iterator<Item = &'genv fhir::PrimOpProp<'genv>> {
599        self.fhir_crate().items.values().filter_map(|item| {
600            if let fhir::FluxItem::PrimOpProp(prop) = item { Some(*prop) } else { None }
601        })
602    }
603
604    pub fn fhir_get_generics(
605        self,
606        def_id: LocalDefId,
607    ) -> QueryResult<Option<&'genv fhir::Generics<'genv>>> {
608        if matches!(self.def_kind(def_id), DefKind::Closure) {
610            Ok(None)
611        } else {
612            Ok(Some(self.fhir_expect_owner_node(def_id)?.generics()))
613        }
614    }
615
616    pub fn fhir_expect_refinement_kind(
617        self,
618        def_id: LocalDefId,
619    ) -> QueryResult<&'genv fhir::RefinementKind<'genv>> {
620        let kind = match &self.fhir_expect_item(def_id)?.kind {
621            fhir::ItemKind::Enum(enum_def) => &enum_def.refinement,
622            fhir::ItemKind::Struct(struct_def) => &struct_def.refinement,
623            _ => bug!("expected struct, enum or type alias"),
624        };
625        Ok(kind)
626    }
627
628    pub fn fhir_expect_item(self, def_id: LocalDefId) -> QueryResult<&'genv fhir::Item<'genv>> {
629        if let fhir::Node::Item(item) = self.fhir_node(def_id)? {
630            Ok(item)
631        } else {
632            Err(query_bug!(def_id, "expected item: `{def_id:?}`"))
633        }
634    }
635
636    pub fn fhir_expect_owner_node(self, def_id: LocalDefId) -> QueryResult<fhir::OwnerNode<'genv>> {
637        let Some(owner) = self.fhir_node(def_id)?.as_owner() else {
638            return Err(query_bug!(def_id, "cannot find owner node"));
639        };
640        Ok(owner)
641    }
642
643    pub fn fhir_node(self, def_id: LocalDefId) -> QueryResult<fhir::Node<'genv>> {
644        self.desugar(def_id)
645    }
646}
647
648#[macro_export]
649macro_rules! try_alloc_slice {
650    ($genv:expr, $slice:expr, $map:expr $(,)?) => {{
651        let slice = $slice;
652        $crate::try_alloc_slice!($genv, cap: slice.len(), slice.into_iter().map($map))
653    }};
654    ($genv:expr, cap: $cap:expr, $it:expr $(,)?) => {{
655        let mut err = None;
656        let slice = $genv.alloc_slice_with_capacity($cap, $it.into_iter().collect_errors(&mut err));
657        err.map_or(Ok(slice), Err)
658    }};
659}
660
661impl ErrorEmitter for GlobalEnv<'_, '_> {
662    fn emit<'a>(&'a self, err: impl rustc_errors::Diagnostic<'a>) -> rustc_span::ErrorGuaranteed {
663        self.sess().emit(err)
664    }
665}