1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134

use flux_arc_interner::List;
use flux_common::{bug, tracked_span_bug};
use rustc_hir::def::DefKind;
use rustc_span::def_id::DefId;

use crate::{global_env::GlobalEnv, queries::QueryResult, query_bug, rty};

impl<'sess, 'tcx> GlobalEnv<'sess, 'tcx> {
    pub fn sort_of_self_ty_alias(self, alias_to: DefId) -> QueryResult<Option<rty::Sort>> {
        let self_ty = self.tcx().type_of(alias_to).instantiate_identity();
        self.sort_of_rust_ty(alias_to, self_ty)
    }

    pub fn sort_of_generic_param(self, def_id: DefId) -> QueryResult<Option<rty::Sort>> {
        let parent = self.tcx().parent(def_id);
        let index = self.def_id_to_param_index(def_id);
        let param = self.generics_of(parent)?.param_at(index as usize, self)?;
        let sort = match &param.kind {
            rty::GenericParamDefKind::Base { .. } => {
                Some(rty::Sort::Param(rty::ParamTy { index, name: param.name }))
            }
            rty::GenericParamDefKind::Const { .. } => {
                let ty = self.tcx().type_of(def_id).instantiate_identity();
                self.sort_of_rust_ty(parent, ty)?
            }
            rty::GenericParamDefKind::Type { .. } | rty::GenericParamDefKind::Lifetime => None,
        };
        Ok(sort)
    }

    fn sort_of_rust_ty(
        self,
        def_id: DefId,
        ty: rustc_middle::ty::Ty,
    ) -> QueryResult<Option<rty::Sort>> {
        use rustc_middle::ty;
        let sort = match ty.kind() {
            ty::TyKind::Bool => Some(rty::Sort::Bool),
            ty::TyKind::Slice(_) | ty::TyKind::Int(_) | ty::TyKind::Uint(_) => Some(rty::Sort::Int),
            ty::TyKind::Char => Some(rty::Sort::Char),
            ty::TyKind::Str => Some(rty::Sort::Str),
            ty::TyKind::Adt(adt_def, args) => {
                let mut sort_args = vec![];
                let sort_def = self.adt_sort_def_of(adt_def.did())?;
                for arg in sort_def.filter_generic_args(args) {
                    let Some(sort) = self.sort_of_rust_ty(def_id, arg.expect_ty())? else {
                        return Ok(None);
                    };
                    sort_args.push(sort);
                }
                let ctor = rty::SortCtor::Adt(self.adt_sort_def_of(adt_def.did())?);
                Some(rty::Sort::App(ctor, List::from_vec(sort_args)))
            }
            ty::TyKind::Param(p) => {
                let generic_param_def = self
                    .tcx()
                    .generics_of(def_id)
                    .param_at(p.index as usize, self.tcx());
                self.sort_of_generic_param(generic_param_def.def_id)?
            }
            ty::TyKind::Float(_)
            | ty::TyKind::RawPtr(..)
            | ty::TyKind::Ref(..)
            | ty::TyKind::Tuple(_)
            | ty::TyKind::Array(..)
            | ty::TyKind::Alias(..)
            | ty::TyKind::Never => Some(rty::Sort::unit()),
            _ => bug!("unexpected self ty {ty:?}"),
        };
        Ok(sort)
    }

    pub fn normalize_weak_alias_sort(self, alias_ty: &rty::AliasTy) -> QueryResult<rty::Sort> {
        match self.def_kind(alias_ty.def_id) {
            DefKind::Impl { .. } => Ok(self.sort_of_self_ty_alias(alias_ty.def_id)?.unwrap()),
            DefKind::TyAlias => {
                Ok(self
                    .type_of(alias_ty.def_id)?
                    .instantiate(self.tcx(), &alias_ty.args, &alias_ty.refine_args)
                    .expect_ctor()
                    .sort())
            }
            DefKind::Struct | DefKind::Enum => {
                Ok(self
                    .adt_sort_def_of(alias_ty.def_id)?
                    .to_sort(&alias_ty.args))
            }
            _ => Err(query_bug!(alias_ty.def_id, "unexpected weak alias `{:?}`", alias_ty.def_id)),
        }
    }
}

impl rty::BaseTy {
    pub fn sort(&self) -> rty::Sort {
        match self {
            rty::BaseTy::Int(_) | rty::BaseTy::Uint(_) | rty::BaseTy::Slice(_) => rty::Sort::Int,
            rty::BaseTy::Bool => rty::Sort::Bool,
            rty::BaseTy::Char => rty::Sort::Char,
            rty::BaseTy::Adt(adt_def, args) => adt_def.sort(args),
            rty::BaseTy::Param(param_ty) => rty::Sort::Param(*param_ty),
            rty::BaseTy::Str => rty::Sort::Str,
            rty::BaseTy::Alias(kind, alias_ty) => {
                // HACK(nilehmann) The refinement arguments in `alias_ty` should not influence the
                // sort. However, we must explicitly remove them because they can contain expression
                // holes. If we don't remove them, we would generate inference variables for them
                // which we won't be able to solve.
                let alias_ty =
                    rty::AliasTy::new(alias_ty.def_id, alias_ty.args.clone(), List::empty());
                rty::Sort::Alias(*kind, alias_ty)
            }
            rty::BaseTy::Float(_)
            | rty::BaseTy::RawPtr(..)
            | rty::BaseTy::Ref(..)
            | rty::BaseTy::FnPtr(..)
            | rty::BaseTy::FnDef(..)
            | rty::BaseTy::Tuple(_)
            | rty::BaseTy::Array(_, _)
            | rty::BaseTy::Closure(..)
            | rty::BaseTy::Coroutine(..)
            | rty::BaseTy::Dynamic(_, _)
            | rty::BaseTy::Never => rty::Sort::unit(),
            rty::BaseTy::Infer(_) => tracked_span_bug!(),
        }
    }
}

impl rty::AliasReft {
    pub fn fsort(&self, genv: GlobalEnv) -> QueryResult<Option<rty::FuncSort>> {
        let Some(fsort) = genv.sort_of_assoc_reft(self.trait_id, self.name)? else {
            return Ok(None);
        };
        Ok(Some(fsort.instantiate(genv.tcx(), &self.args, &[])))
    }
}