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