flux_middle/
sort_of.rs

1use flux_arc_interner::List;
2use flux_common::tracked_span_bug;
3use rustc_hir::def::DefKind;
4use rustc_span::def_id::DefId;
5
6use crate::{
7    global_env::GlobalEnv,
8    queries::{QueryErr, QueryResult},
9    query_bug,
10    rty::{
11        self,
12        fold::{FallibleTypeFolder, TypeFoldable, TypeSuperFoldable as _},
13    },
14};
15
16impl GlobalEnv<'_, '_> {
17    pub fn sort_of_self_ty_alias(self, alias_to: DefId) -> QueryResult<Option<rty::Sort>> {
18        let self_ty = self.tcx().type_of(alias_to).instantiate_identity();
19        self.sort_of_rust_ty(alias_to, self_ty)
20    }
21
22    pub fn sort_of_generic_param(self, def_id: DefId) -> QueryResult<Option<rty::Sort>> {
23        let parent = self.tcx().parent(def_id);
24        let index = self.def_id_to_param_index(def_id);
25        let param = self.generics_of(parent)?.param_at(index as usize, self)?;
26        let sort = match &param.kind {
27            rty::GenericParamDefKind::Base { .. } => {
28                Some(rty::Sort::Param(rty::ParamTy { index, name: param.name }))
29            }
30            rty::GenericParamDefKind::Const { .. } => {
31                let ty = self.tcx().type_of(def_id).instantiate_identity();
32                self.sort_of_rust_ty(parent, ty)?
33            }
34            rty::GenericParamDefKind::Type { .. } | rty::GenericParamDefKind::Lifetime => None,
35        };
36        Ok(sort)
37    }
38
39    pub fn sort_of_def_id(self, def_id: DefId) -> QueryResult<Option<rty::Sort>> {
40        let ty = self.tcx().type_of(def_id).no_bound_vars().unwrap();
41        if ty.is_integral() { Ok(Some(rty::Sort::Int)) } else { self.sort_of_rust_ty(def_id, ty) }
42    }
43
44    pub fn sort_of_rust_ty(
45        self,
46        def_id: DefId,
47        ty: rustc_middle::ty::Ty,
48    ) -> QueryResult<Option<rty::Sort>> {
49        use rustc_middle::ty;
50        let sort = match ty.kind() {
51            ty::TyKind::Bool => Some(rty::Sort::Bool),
52            ty::TyKind::Slice(_) | ty::TyKind::Int(_) | ty::TyKind::Uint(_) => Some(rty::Sort::Int),
53            ty::TyKind::Char => Some(rty::Sort::Char),
54            ty::TyKind::Str => Some(rty::Sort::Str),
55            ty::TyKind::Adt(adt_def, args) => {
56                let mut sort_args = vec![];
57                let sort_def = self.adt_sort_def_of(adt_def.did())?;
58                for arg in sort_def.filter_generic_args(args) {
59                    let Some(sort) = self.sort_of_rust_ty(def_id, arg.expect_ty())? else {
60                        return Ok(None);
61                    };
62                    sort_args.push(sort);
63                }
64                let ctor = rty::SortCtor::Adt(self.adt_sort_def_of(adt_def.did())?);
65                Some(rty::Sort::App(ctor, List::from_vec(sort_args)))
66            }
67            ty::TyKind::Param(p) => {
68                let generic_param_def = self
69                    .tcx()
70                    .generics_of(def_id)
71                    .param_at(p.index as usize, self.tcx());
72                self.sort_of_generic_param(generic_param_def.def_id)?
73            }
74            ty::TyKind::Float(_)
75            | ty::TyKind::RawPtr(..)
76            | ty::TyKind::Ref(..)
77            | ty::TyKind::Tuple(_)
78            | ty::TyKind::Array(..)
79            | ty::TyKind::Alias(..)
80            | ty::TyKind::Never => Some(rty::Sort::unit()),
81            _ => None,
82        };
83        Ok(sort)
84    }
85
86    pub fn normalize_weak_alias_sort(self, alias_ty: &rty::AliasTy) -> QueryResult<rty::Sort> {
87        match self.def_kind(alias_ty.def_id) {
88            DefKind::Impl { .. } => Ok(self.sort_of_self_ty_alias(alias_ty.def_id)?.unwrap()),
89            DefKind::TyAlias => {
90                Ok(self
91                    .type_of(alias_ty.def_id)?
92                    .instantiate(self.tcx(), &alias_ty.args, &alias_ty.refine_args)
93                    .expect_ctor()
94                    .sort())
95            }
96            DefKind::Struct | DefKind::Enum => {
97                Ok(self
98                    .adt_sort_def_of(alias_ty.def_id)?
99                    .to_sort(&alias_ty.args))
100            }
101            _ => Err(query_bug!(alias_ty.def_id, "unexpected weak alias `{:?}`", alias_ty.def_id)),
102        }
103    }
104
105    pub fn deep_normalize_weak_alias_sorts<T: TypeFoldable>(self, t: &T) -> QueryResult<T> {
106        struct WeakAliasSortNormalizer<'genv, 'tcx> {
107            genv: GlobalEnv<'genv, 'tcx>,
108        }
109
110        impl FallibleTypeFolder for WeakAliasSortNormalizer<'_, '_> {
111            type Error = QueryErr;
112
113            fn try_fold_sort(&mut self, sort: &rty::Sort) -> QueryResult<rty::Sort> {
114                if let rty::Sort::Alias(rty::AliasKind::Weak, alias_ty) = sort {
115                    self.genv
116                        .normalize_weak_alias_sort(alias_ty)?
117                        .try_fold_with(self)
118                } else {
119                    sort.try_super_fold_with(self)
120                }
121            }
122        }
123        t.try_fold_with(&mut WeakAliasSortNormalizer { genv: self })
124    }
125}
126
127impl rty::BaseTy {
128    pub fn sort(&self) -> rty::Sort {
129        match self {
130            rty::BaseTy::Int(_) | rty::BaseTy::Uint(_) | rty::BaseTy::Slice(_) => rty::Sort::Int,
131            rty::BaseTy::Bool => rty::Sort::Bool,
132            rty::BaseTy::Char => rty::Sort::Char,
133            rty::BaseTy::Adt(adt_def, args) => adt_def.sort(args),
134            rty::BaseTy::Param(param_ty) => rty::Sort::Param(*param_ty),
135            rty::BaseTy::Str => rty::Sort::Str,
136            rty::BaseTy::Alias(kind, alias_ty) => {
137                // HACK(nilehmann) The refinement arguments in `alias_ty` should not influence the
138                // sort. However, we must explicitly remove them because they can contain expression
139                // holes. If we don't remove them, we would generate inference variables for them
140                // which we won't be able to solve.
141                let alias_ty =
142                    rty::AliasTy::new(alias_ty.def_id, alias_ty.args.clone(), List::empty());
143                rty::Sort::Alias(*kind, alias_ty)
144            }
145            rty::BaseTy::Float(_)
146            | rty::BaseTy::RawPtr(..)
147            | rty::BaseTy::RawPtrMetadata(..) // TODO(RJ): This should be `int` for slice?
148            | rty::BaseTy::Ref(..)
149            | rty::BaseTy::FnPtr(..)
150            | rty::BaseTy::FnDef(..)
151            | rty::BaseTy::Tuple(_)
152            | rty::BaseTy::Array(_, _)
153            | rty::BaseTy::Closure(..)
154            | rty::BaseTy::Coroutine(..)
155            | rty::BaseTy::Dynamic(_, _)
156            | rty::BaseTy::Never
157            | rty::BaseTy::Foreign(..) => rty::Sort::unit(),
158            rty::BaseTy::Infer(_) => tracked_span_bug!(),
159        }
160    }
161}
162
163impl rty::AliasReft {
164    pub fn fsort(&self, genv: GlobalEnv) -> QueryResult<rty::FuncSort> {
165        Ok(genv
166            .sort_of_assoc_reft(self.assoc_id)?
167            .instantiate(genv.tcx(), &self.args, &[]))
168    }
169}