Skip to main content

flux_middle/
builtin_assoc_refts.rs

1use std::sync::OnceLock;
2
3use flux_arc_interner::List;
4use flux_common::bug;
5use flux_syntax::symbols::sym;
6use rustc_data_structures::unord::UnordMap;
7use rustc_hir::{LangItem, def_id::DefId};
8use rustc_span::{DUMMY_SP, Symbol};
9
10use crate::{
11    def_id::FluxDefId,
12    global_env::GlobalEnv,
13    rty::{self, AliasReft, AssocRefinements, AssocReft, BaseTy},
14};
15
16impl<'tcx> GlobalEnv<'_, 'tcx> {
17    #[allow(
18        clippy::disallowed_methods,
19        reason = "this file is the source of truth for builtin assoc refts def ids"
20    )]
21    pub fn builtin_assoc_refts(self, def_id: DefId) -> Option<AssocRefinements> {
22        static BUILTIN: OnceLock<UnordMap<DefId, AssocRefinements>> = OnceLock::new();
23
24        BUILTIN
25            .get_or_init(|| {
26                let tcx = self.tcx();
27
28                let mut map = UnordMap::default();
29
30                // FnOnce
31                let fn_once_id = tcx.require_lang_item(LangItem::FnOnce, DUMMY_SP);
32                map.insert(
33                    fn_once_id,
34                    AssocRefinements {
35                        items: List::from_arr([AssocReft::new(
36                            FluxDefId::new(fn_once_id, sym::no_panic),
37                            false,
38                            tcx.def_span(fn_once_id),
39                        )]),
40                    },
41                );
42
43                // Sized
44                let sized_id = tcx.require_lang_item(LangItem::Sized, DUMMY_SP);
45                map.insert(
46                    sized_id,
47                    AssocRefinements {
48                        items: List::from_arr([
49                            AssocReft::new(
50                                FluxDefId::new(sized_id, sym::size_of),
51                                false,
52                                tcx.def_span(sized_id),
53                            ),
54                            AssocReft::new(
55                                FluxDefId::new(sized_id, sym::align_of),
56                                false,
57                                tcx.def_span(sized_id),
58                            ),
59                        ]),
60                    },
61                );
62                map
63            })
64            .get(&def_id)
65            .cloned()
66    }
67
68    #[allow(
69        clippy::disallowed_methods,
70        reason = "this file is the source of truth for builtin assoc refts def ids"
71    )]
72    pub fn builtin_assoc_reft_sort(
73        self,
74        assoc_id: FluxDefId,
75    ) -> Option<rty::EarlyBinder<rty::FuncSort>> {
76        static BUILTIN: OnceLock<UnordMap<FluxDefId, rty::FuncSort>> = OnceLock::new();
77
78        BUILTIN
79            .get_or_init(|| {
80                let tcx = self.tcx();
81
82                let mut map = UnordMap::default();
83
84                // FnOnce
85                let fn_once_id = tcx.require_lang_item(LangItem::FnOnce, DUMMY_SP);
86                map.insert(
87                    FluxDefId::new(fn_once_id, sym::no_panic),
88                    rty::FuncSort::new(vec![], rty::Sort::Bool),
89                );
90
91                // Sized
92                let sized_id = tcx.require_lang_item(LangItem::Sized, DUMMY_SP);
93                map.insert(
94                    FluxDefId::new(sized_id, sym::size_of),
95                    rty::FuncSort::new(vec![], rty::Sort::Int),
96                );
97                map.insert(
98                    FluxDefId::new(sized_id, sym::align_of),
99                    rty::FuncSort::new(vec![], rty::Sort::Int),
100                );
101                map
102            })
103            .get(&assoc_id)
104            .cloned()
105            .map(rty::EarlyBinder)
106    }
107
108    /// Returns the [`FluxDefId`] for a built-in associated refinement given the parent trait's
109    /// [`DefId`] and the assoc reft [`Symbol`]. Panics if no such built-in assoc reft exists.
110    ///
111    /// This is the single source of truth for looking up built-in assoc reft ids; callers should
112    /// use this instead of constructing a [`FluxDefId`] directly.
113    pub fn require_builtin_assoc_reft(self, def_id: DefId, name: Symbol) -> FluxDefId {
114        self.builtin_assoc_refts(def_id)
115            .and_then(|assoc_refts| {
116                assoc_refts
117                    .items
118                    .iter()
119                    .find(|reft| reft.name() == name)
120                    .map(|reft| reft.def_id())
121            })
122            .unwrap_or_else(|| bug!("missing builtin assoc reft `{name}` for `{def_id:?}`"))
123    }
124
125    pub fn builtin_assoc_reft_body(
126        self,
127        typing_env: rustc_middle::ty::TypingEnv<'tcx>,
128        alias_reft: &AliasReft,
129    ) -> rty::Lambda {
130        let tcx = self.tcx();
131
132        if tcx.is_lang_item(alias_reft.assoc_id.parent(), LangItem::Sized) {
133            let self_ty = alias_reft.to_rustc_trait_ref(tcx).self_ty();
134            let layout = tcx.layout_of(typing_env.as_query_input(self_ty)).unwrap();
135            let body = match alias_reft.assoc_id.name() {
136                sym::size_of => rty::Expr::constant(rty::Constant::from(layout.size.bytes())),
137                sym::align_of => rty::Expr::constant(rty::Constant::from(layout.align.abi.bytes())),
138                _ => bug!("invalid builtin assoc reft {:?}", alias_reft.assoc_id),
139            };
140            rty::Lambda::bind_with_vars(body, List::empty(), rty::Sort::Int)
141        } else if tcx.is_lang_item(alias_reft.assoc_id.parent(), LangItem::FnOnce)
142            && alias_reft.assoc_id.name() == sym::no_panic
143        {
144            let self_ty = alias_reft.self_ty();
145            let body = match self_ty.as_bty_skipping_binder() {
146                BaseTy::Closure(_, _, _, no_panic) => {
147                    if *no_panic {
148                        rty::Expr::tt()
149                    } else {
150                        rty::Expr::ff()
151                    }
152                }
153                _ => rty::Expr::ff(),
154            };
155            rty::Lambda::bind_with_vars(body, List::empty(), rty::Sort::Bool)
156        } else {
157            bug!("invalid builtin assoc reft {:?}", alias_reft.assoc_id)
158        }
159    }
160}