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