flux_driver/collector/
detached_specs.rs

1use std::collections::{HashMap, hash_map::Entry};
2
3use flux_common::dbg::{self, SpanTrace};
4use flux_syntax::surface::{self, DetachedItem, ExprPath, NodeId};
5use itertools::Itertools;
6use rustc_errors::ErrorGuaranteed;
7use rustc_hir::{
8    OwnerId,
9    def::{DefKind, Res},
10    def_id::LocalDefId,
11};
12use rustc_middle::ty::{AssocItem, AssocKind, Ty, TyCtxt};
13use rustc_span::{Symbol, def_id::DefId};
14
15use crate::collector::{FluxAttrs, SpecCollector, errors};
16type Result<T = ()> = std::result::Result<T, ErrorGuaranteed>;
17
18#[derive(PartialEq, Eq, Debug, Hash, Clone, Copy)]
19enum LookupRes {
20    DefId(DefId),
21    Name(Symbol),
22}
23
24impl LookupRes {
25    fn from_name<T: std::fmt::Debug>(thing: &T) -> Self {
26        let str = format!("{thing:?}");
27        LookupRes::Name(Symbol::intern(&str))
28    }
29
30    fn new(ty: &Ty) -> Self {
31        match ty.kind() {
32            rustc_middle::ty::TyKind::Adt(adt_def, _) => LookupRes::DefId(adt_def.did()),
33            _ => Self::from_name(ty),
34        }
35    }
36}
37
38#[derive(PartialEq, Eq, Debug, Hash)]
39struct TraitImplKey {
40    trait_: LookupRes,
41    self_ty: LookupRes,
42}
43
44fn path_to_symbol(path: &surface::ExprPath) -> Symbol {
45    let path_string = format!(
46        "{}",
47        path.segments
48            .iter()
49            .format_with("::", |s, f| f(&s.ident.name))
50    );
51    Symbol::intern(&path_string)
52}
53
54fn item_def_kind(kind: &surface::DetachedItemKind) -> Vec<DefKind> {
55    match kind {
56        surface::DetachedItemKind::FnSig(_) => vec![DefKind::Fn],
57        surface::DetachedItemKind::Mod(_) => vec![DefKind::Mod],
58        surface::DetachedItemKind::Struct(_) => vec![DefKind::Struct],
59        surface::DetachedItemKind::Enum(_) => vec![DefKind::Enum],
60        surface::DetachedItemKind::InherentImpl(_) | surface::DetachedItemKind::TraitImpl(_) => {
61            vec![DefKind::Struct, DefKind::Enum]
62        }
63        surface::DetachedItemKind::Trait(_) => vec![DefKind::Trait],
64        surface::DetachedItemKind::Static(_) => {
65            vec![DefKind::Static {
66                mutability: rustc_ast::Mutability::Not,
67                nested: false,
68                safety: rustc_hir::Safety::Safe,
69            }]
70        }
71    }
72}
73
74#[derive(Debug)]
75struct ScopeResolver {
76    items: HashMap<(Symbol, DefKind), LookupRes>,
77}
78
79impl ScopeResolver {
80    fn new(tcx: TyCtxt, def_id: LocalDefId, impl_resolver: &TraitImplResolver) -> Self {
81        let mut items = HashMap::default();
82        for child in tcx.module_children_local(def_id) {
83            let ident = child.ident;
84            if let Res::Def(exp_kind, def_id) = child.res {
85                items.insert((ident.name, exp_kind), LookupRes::DefId(def_id));
86            }
87        }
88        for pty in rustc_hir::PrimTy::ALL {
89            let name = pty.name();
90            items.insert((name, DefKind::Struct), LookupRes::Name(name)); // HACK: use DefKind::Struct for primitive...
91        }
92        for trait_impl_key in impl_resolver.items.keys() {
93            if let LookupRes::DefId(trait_id) = trait_impl_key.trait_ {
94                let name = Symbol::intern(&tcx.def_path_str(trait_id));
95                items.insert((name, DefKind::Trait), trait_impl_key.trait_);
96            }
97        }
98        Self { items }
99    }
100
101    fn lookup(&self, path: &ExprPath, item_kind: &surface::DetachedItemKind) -> Option<LookupRes> {
102        let symbol = path_to_symbol(path);
103        for kind in item_def_kind(item_kind) {
104            let key = (symbol, kind);
105            if let Some(res) = self.items.get(&key) {
106                return Some(*res);
107            }
108        }
109        None
110    }
111}
112
113#[derive(Debug)]
114struct TraitImplResolver {
115    items: HashMap<TraitImplKey, LocalDefId>,
116}
117
118impl TraitImplResolver {
119    fn new(tcx: TyCtxt) -> Self {
120        let mut items = HashMap::default();
121        for (trait_id, impl_ids) in tcx.all_local_trait_impls(()) {
122            let trait_ = LookupRes::DefId(*trait_id);
123            for impl_id in impl_ids {
124                let poly_trait_ref = tcx.impl_trait_ref(*impl_id);
125                let self_ty = poly_trait_ref
126                    .instantiate_identity()
127                    .skip_norm_wip()
128                    .self_ty();
129                let self_ty = LookupRes::new(&self_ty);
130                let key = TraitImplKey { trait_, self_ty };
131                items.insert(key, *impl_id);
132            }
133        }
134        Self { items }
135    }
136
137    fn resolve(&self, trait_: LookupRes, self_ty: LookupRes) -> Option<LocalDefId> {
138        let key = TraitImplKey { trait_, self_ty };
139        self.items.get(&key).copied()
140    }
141}
142
143pub(super) struct DetachedSpecsCollector<'a, 'sess, 'tcx> {
144    inner: &'a mut SpecCollector<'sess, 'tcx>,
145    id_resolver: HashMap<NodeId, LookupRes>,
146    impl_resolver: TraitImplResolver,
147}
148
149impl<'a, 'sess, 'tcx> DetachedSpecsCollector<'a, 'sess, 'tcx> {
150    pub(super) fn collect(
151        inner: &'a mut SpecCollector<'sess, 'tcx>,
152        attrs: &mut FluxAttrs,
153        module_id: LocalDefId,
154    ) -> Result {
155        if let Some(detached_specs) = attrs.detached_specs() {
156            let trait_impl_resolver = TraitImplResolver::new(inner.tcx);
157            let mut collector =
158                Self { inner, id_resolver: HashMap::default(), impl_resolver: trait_impl_resolver };
159            collector.run(detached_specs, module_id)?;
160        };
161        Ok(())
162    }
163
164    fn run(&mut self, detached_specs: surface::DetachedSpecs, def_id: LocalDefId) -> Result {
165        self.resolve(&detached_specs, def_id)?;
166        for item in detached_specs.items {
167            self.attach(item)?;
168        }
169        Ok(())
170    }
171
172    fn resolve_path_kind(
173        &mut self,
174        resolver: &ScopeResolver,
175        path: &ExprPath,
176        kind: &surface::DetachedItemKind,
177    ) -> Result {
178        let Some(res) = resolver.lookup(path, kind) else {
179            return Err(self
180                .inner
181                .errors
182                .emit(errors::UnresolvedSpecification::new(path, "name")));
183        };
184        self.id_resolver.insert(path.node_id, res);
185        Ok(())
186    }
187
188    fn resolve(&mut self, detached_specs: &surface::DetachedSpecs, def_id: LocalDefId) -> Result {
189        let resolver = ScopeResolver::new(self.inner.tcx, def_id, &self.impl_resolver);
190        for item in &detached_specs.items {
191            self.resolve_path_kind(&resolver, &item.path, &item.kind)?;
192            if let surface::DetachedItemKind::TraitImpl(trait_impl) = &item.kind {
193                let kind = surface::DetachedItemKind::Trait(surface::DetachedTrait::default());
194                self.resolve_path_kind(&resolver, &trait_impl.trait_, &kind)?;
195            }
196        }
197        Ok(())
198    }
199
200    #[allow(
201        clippy::disallowed_methods,
202        reason = "this is pre-extern specs so it's fine: https://flux-rs.zulipchat.com/#narrow/channel/486369-verify-std/topic/detached-specs/near/529548357"
203    )]
204    fn unwrap_def_id(&self, def_id: &DefId) -> Result<Option<LocalDefId>> {
205        Ok(def_id.as_local())
206    }
207
208    fn lookup(&mut self, item: &surface::DetachedItem) -> Result<LocalDefId> {
209        let path_def_id = self.id_resolver.get(&item.path.node_id);
210
211        if let surface::DetachedItemKind::TraitImpl(trait_impl) = &item.kind
212            && let Some(trait_) = self.id_resolver.get(&trait_impl.trait_.node_id)
213            && let Some(self_ty) = path_def_id
214            && let Some(impl_id) = self.impl_resolver.resolve(*trait_, *self_ty)
215        {
216            return Ok(impl_id);
217        }
218        if let Some(LookupRes::DefId(def_id)) = self.id_resolver.get(&item.path.node_id)
219            && let Some(local_def_id) = self.unwrap_def_id(def_id)?
220        {
221            return Ok(local_def_id);
222        }
223        Err(self
224            .inner
225            .errors
226            .emit(errors::UnresolvedSpecification::new(&item.path, "item")))
227    }
228
229    fn attach(&mut self, item: surface::DetachedItem) -> Result {
230        let def_id = self.lookup(&item)?;
231        let owner_id = self.inner.tcx.local_def_id_to_hir_id(def_id).owner;
232        let span = item.span();
233        let dst_span = self.inner.tcx.def_span(def_id);
234        dbg::hyperlink!(self.inner.tcx, span, dst_span);
235        match item.kind {
236            surface::DetachedItemKind::FnSig(fn_sig) => {
237                self.inner.insert_item(
238                    owner_id,
239                    surface::Item {
240                        attrs: item.attrs,
241                        kind: surface::ItemKind::Fn(Some(fn_sig)),
242                        node_id: item.node_id,
243                    },
244                )?;
245            }
246            surface::DetachedItemKind::Struct(struct_def) => {
247                self.inner.insert_item(
248                    owner_id,
249                    surface::Item {
250                        attrs: item.attrs,
251                        kind: surface::ItemKind::Struct(struct_def),
252                        node_id: item.node_id,
253                    },
254                )?;
255            }
256            surface::DetachedItemKind::Enum(enum_def) => {
257                self.inner.insert_item(
258                    owner_id,
259                    surface::Item {
260                        attrs: item.attrs,
261                        kind: surface::ItemKind::Enum(enum_def),
262                        node_id: item.node_id,
263                    },
264                )?;
265            }
266            surface::DetachedItemKind::Mod(detached_specs) => {
267                self.run(detached_specs, owner_id.def_id)?;
268            }
269            surface::DetachedItemKind::Trait(trait_def) => {
270                self.collect_trait(owner_id, item.node_id, item.attrs, trait_def)?;
271            }
272            surface::DetachedItemKind::InherentImpl(inherent_impl) => {
273                let tcx = self.inner.tcx;
274                let assoc_items = tcx
275                    .inherent_impls(def_id)
276                    .iter()
277                    .flat_map(|impl_id| tcx.associated_items(*impl_id).in_definition_order());
278                self.collect_assoc_methods(
279                    inherent_impl.items,
280                    assoc_items,
281                    |this, owner_id, item| {
282                        this.inner.insert_impl_item(
283                            owner_id,
284                            surface::ImplItemFn {
285                                attrs: item.attrs,
286                                sig: Some(item.kind),
287                                node_id: item.node_id,
288                            },
289                        )
290                    },
291                )?;
292            }
293            surface::DetachedItemKind::TraitImpl(trait_impl) => {
294                self.collect_trait_impl(owner_id, item.node_id, item.attrs, trait_impl)?;
295            }
296            surface::DetachedItemKind::Static(static_info) => {
297                self.inner.insert_item(
298                    owner_id,
299                    surface::Item {
300                        attrs: item.attrs,
301                        kind: surface::ItemKind::Static(static_info),
302                        node_id: item.node_id,
303                    },
304                )?;
305            }
306        };
307        Ok(())
308    }
309
310    fn collect_trait(
311        &mut self,
312        owner_id: OwnerId,
313        node_id: NodeId,
314        attrs: Vec<surface::Attr>,
315        trait_def: surface::DetachedTrait,
316    ) -> Result {
317        // 1. Collect the associated-refinements
318        self.inner.insert_item(
319            owner_id,
320            surface::Item {
321                attrs,
322                kind: surface::ItemKind::Trait(surface::Trait {
323                    generics: None,
324                    assoc_refinements: trait_def.refts,
325                }),
326                node_id,
327            },
328        )?;
329
330        // 2. Collect the method specifications
331        let tcx = self.inner.tcx;
332        let assoc_items = tcx.associated_items(owner_id.def_id).in_definition_order();
333        self.collect_assoc_methods(trait_def.items, assoc_items, |this, owner_id, item| {
334            this.inner.insert_trait_item(
335                owner_id,
336                surface::TraitItemFn {
337                    attrs: item.attrs,
338                    sig: Some(item.kind),
339                    node_id: item.node_id,
340                },
341            )
342        })
343    }
344
345    fn collect_trait_impl(
346        &mut self,
347        owner_id: OwnerId,
348        node_id: NodeId,
349        attrs: Vec<surface::Attr>,
350        trait_impl: surface::DetachedTraitImpl,
351    ) -> Result {
352        // 1. Collect the associated-refinements
353        self.inner.insert_item(
354            owner_id,
355            surface::Item {
356                attrs,
357                kind: surface::ItemKind::Impl(surface::Impl {
358                    generics: None,
359                    assoc_refinements: trait_impl.refts,
360                }),
361                node_id,
362            },
363        )?;
364
365        // 2. Collect the method specifications
366        let tcx = self.inner.tcx;
367        let assoc_items = tcx.associated_items(owner_id.def_id).in_definition_order();
368        self.collect_assoc_methods(trait_impl.items, assoc_items, |this, owner_id, item| {
369            this.inner.insert_impl_item(
370                owner_id,
371                surface::ImplItemFn {
372                    attrs: item.attrs,
373                    sig: Some(item.kind),
374                    node_id: item.node_id,
375                },
376            )
377        })
378    }
379
380    fn collect_assoc_methods(
381        &mut self,
382        methods: Vec<DetachedItem<surface::FnSig>>,
383        assoc_items: impl Iterator<Item = &'tcx AssocItem>,
384        mut insert_item: impl FnMut(&mut Self, OwnerId, DetachedItem<surface::FnSig>) -> Result,
385    ) -> Result {
386        let mut table: HashMap<Symbol, DetachedItem<(surface::FnSig, Option<DefId>)>> =
387            HashMap::default();
388        // 1. make a table of the impl-items
389        for item in methods {
390            let name = path_to_symbol(&item.path);
391            let span = item.path.span;
392            if let Entry::Occupied(_) = table.entry(name) {
393                return Err(self
394                    .inner
395                    .errors
396                    .emit(errors::MultipleSpecifications { name, span }));
397            } else {
398                table.insert(name, item.map_kind(|spec| (spec, None)));
399            }
400        }
401        // 2. walk over all the assoc-items to resolve names
402        for item in assoc_items {
403            if let AssocKind::Fn { name, .. } = item.kind
404                && let Some(val) = table.get_mut(&name)
405                && val.kind.1.is_none()
406            {
407                val.kind.1 = Some(item.def_id);
408            }
409        }
410        // 3. Attach the `fn_sig` to the resolved `DefId`
411        for (_name, item) in table {
412            let Some(def_id) = item.kind.1 else {
413                return Err(self
414                    .inner
415                    .errors
416                    .emit(errors::UnresolvedSpecification::new(&item.path, "identifier")));
417            };
418            if let Some(def_id) = self.unwrap_def_id(&def_id)? {
419                dbg::hyperlink!(self.inner.tcx, item.path.span, self.inner.tcx.def_span(def_id));
420                let owner_id = self.inner.tcx.local_def_id_to_hir_id(def_id).owner;
421                insert_item(self, owner_id, item.map_kind(|k| k.0))?;
422            }
423        }
424        Ok(())
425    }
426}
flux_driver/collector/detached_specs.rs

flux_driver/collector/
detached_specs.rs
