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    }
65}
66
67#[derive(Debug)]
68struct ScopeResolver {
69    items: HashMap<(Symbol, DefKind), LookupRes>,
70}
71
72impl ScopeResolver {
73    fn new(tcx: TyCtxt, def_id: LocalDefId, impl_resolver: &TraitImplResolver) -> Self {
74        let mut items = HashMap::default();
75        for child in tcx.module_children_local(def_id) {
76            let ident = child.ident;
77            if let Res::Def(exp_kind, def_id) = child.res {
78                items.insert((ident.name, exp_kind), LookupRes::DefId(def_id));
79            }
80        }
81        for pty in rustc_hir::PrimTy::ALL {
82            let name = pty.name();
83            items.insert((name, DefKind::Struct), LookupRes::Name(name)); // HACK: use DefKind::Struct for primitive...
84        }
85        for trait_impl_key in impl_resolver.items.keys() {
86            if let LookupRes::DefId(trait_id) = trait_impl_key.trait_ {
87                let name = Symbol::intern(&tcx.def_path_str(trait_id));
88                items.insert((name, DefKind::Trait), trait_impl_key.trait_);
89            }
90        }
91        Self { items }
92    }
93
94    fn lookup(&self, path: &ExprPath, item_kind: &surface::DetachedItemKind) -> Option<LookupRes> {
95        let symbol = path_to_symbol(path);
96        for kind in item_def_kind(item_kind) {
97            let key = (symbol, kind);
98            if let Some(res) = self.items.get(&key) {
99                return Some(*res);
100            }
101        }
102        None
103    }
104}
105
106#[derive(Debug)]
107struct TraitImplResolver {
108    items: HashMap<TraitImplKey, LocalDefId>,
109}
110
111impl TraitImplResolver {
112    fn new(tcx: TyCtxt) -> Self {
113        let mut items = HashMap::default();
114        for (trait_id, impl_ids) in tcx.all_local_trait_impls(()) {
115            let trait_ = LookupRes::DefId(*trait_id);
116            for impl_id in impl_ids {
117                if let Some(poly_trait_ref) = tcx.impl_trait_ref(*impl_id) {
118                    let self_ty = poly_trait_ref.instantiate_identity().self_ty();
119                    let self_ty = LookupRes::new(&self_ty);
120                    let key = TraitImplKey { trait_, self_ty };
121                    items.insert(key, *impl_id);
122                }
123            }
124        }
125        Self { items }
126    }
127
128    fn resolve(&self, trait_: LookupRes, self_ty: LookupRes) -> Option<LocalDefId> {
129        let key = TraitImplKey { trait_, self_ty };
130        self.items.get(&key).copied()
131    }
132}
133
134pub(super) struct DetachedSpecsCollector<'a, 'sess, 'tcx> {
135    inner: &'a mut SpecCollector<'sess, 'tcx>,
136    id_resolver: HashMap<NodeId, LookupRes>,
137    impl_resolver: TraitImplResolver,
138}
139
140impl<'a, 'sess, 'tcx> DetachedSpecsCollector<'a, 'sess, 'tcx> {
141    pub(super) fn collect(
142        inner: &'a mut SpecCollector<'sess, 'tcx>,
143        attrs: &mut FluxAttrs,
144        module_id: LocalDefId,
145    ) -> Result {
146        if let Some(detached_specs) = attrs.detached_specs() {
147            let trait_impl_resolver = TraitImplResolver::new(inner.tcx);
148            let mut collector =
149                Self { inner, id_resolver: HashMap::default(), impl_resolver: trait_impl_resolver };
150            collector.run(detached_specs, module_id)?;
151        };
152        Ok(())
153    }
154
155    fn run(&mut self, detached_specs: surface::DetachedSpecs, def_id: LocalDefId) -> Result {
156        self.resolve(&detached_specs, def_id)?;
157        for item in detached_specs.items {
158            self.attach(item)?;
159        }
160        Ok(())
161    }
162
163    fn resolve_path_kind(
164        &mut self,
165        resolver: &ScopeResolver,
166        path: &ExprPath,
167        kind: &surface::DetachedItemKind,
168    ) -> Result {
169        let Some(res) = resolver.lookup(path, kind) else {
170            return Err(self
171                .inner
172                .errors
173                .emit(errors::UnresolvedSpecification::new(path, "name")));
174        };
175        self.id_resolver.insert(path.node_id, res);
176        Ok(())
177    }
178
179    fn resolve(&mut self, detached_specs: &surface::DetachedSpecs, def_id: LocalDefId) -> Result {
180        let resolver = ScopeResolver::new(self.inner.tcx, def_id, &self.impl_resolver);
181        for item in &detached_specs.items {
182            self.resolve_path_kind(&resolver, &item.path, &item.kind)?;
183            if let surface::DetachedItemKind::TraitImpl(trait_impl) = &item.kind {
184                let kind = surface::DetachedItemKind::Trait(surface::DetachedTrait::default());
185                self.resolve_path_kind(&resolver, &trait_impl.trait_, &kind)?;
186            }
187        }
188        Ok(())
189    }
190
191    #[allow(
192        clippy::disallowed_methods,
193        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"
194    )]
195    fn unwrap_def_id(&self, def_id: &DefId) -> Result<Option<LocalDefId>> {
196        Ok(def_id.as_local())
197    }
198
199    fn lookup(&mut self, item: &surface::DetachedItem) -> Result<LocalDefId> {
200        let path_def_id = self.id_resolver.get(&item.path.node_id);
201
202        if let surface::DetachedItemKind::TraitImpl(trait_impl) = &item.kind
203            && let Some(trait_) = self.id_resolver.get(&trait_impl.trait_.node_id)
204            && let Some(self_ty) = path_def_id
205            && let Some(impl_id) = self.impl_resolver.resolve(*trait_, *self_ty)
206        {
207            return Ok(impl_id);
208        }
209        if let Some(LookupRes::DefId(def_id)) = self.id_resolver.get(&item.path.node_id)
210            && let Some(local_def_id) = self.unwrap_def_id(def_id)?
211        {
212            return Ok(local_def_id);
213        }
214        Err(self
215            .inner
216            .errors
217            .emit(errors::UnresolvedSpecification::new(&item.path, "item")))
218    }
219
220    fn attach(&mut self, item: surface::DetachedItem) -> Result {
221        let def_id = self.lookup(&item)?;
222        let owner_id = self.inner.tcx.local_def_id_to_hir_id(def_id).owner;
223        let span = item.span();
224        let dst_span = self.inner.tcx.def_span(def_id);
225        dbg::hyperlink!(self.inner.tcx, span, dst_span);
226        match item.kind {
227            surface::DetachedItemKind::FnSig(fn_sig) => {
228                self.inner.insert_item(
229                    owner_id,
230                    surface::Item {
231                        attrs: item.attrs,
232                        kind: surface::ItemKind::Fn(Some(fn_sig)),
233                        node_id: item.node_id,
234                    },
235                )?;
236            }
237            surface::DetachedItemKind::Struct(struct_def) => {
238                self.inner.insert_item(
239                    owner_id,
240                    surface::Item {
241                        attrs: item.attrs,
242                        kind: surface::ItemKind::Struct(struct_def),
243                        node_id: item.node_id,
244                    },
245                )?;
246            }
247            surface::DetachedItemKind::Enum(enum_def) => {
248                self.inner.insert_item(
249                    owner_id,
250                    surface::Item {
251                        attrs: item.attrs,
252                        kind: surface::ItemKind::Enum(enum_def),
253                        node_id: item.node_id,
254                    },
255                )?;
256            }
257            surface::DetachedItemKind::Mod(detached_specs) => {
258                self.run(detached_specs, owner_id.def_id)?;
259            }
260            surface::DetachedItemKind::Trait(trait_def) => {
261                self.collect_trait(owner_id, item.node_id, item.attrs, trait_def)?;
262            }
263            surface::DetachedItemKind::InherentImpl(inherent_impl) => {
264                let tcx = self.inner.tcx;
265                let assoc_items = tcx
266                    .inherent_impls(def_id)
267                    .iter()
268                    .flat_map(|impl_id| tcx.associated_items(impl_id).in_definition_order());
269                self.collect_assoc_methods(
270                    inherent_impl.items,
271                    assoc_items,
272                    |this, owner_id, item| {
273                        this.inner.insert_impl_item(
274                            owner_id,
275                            surface::ImplItemFn {
276                                attrs: item.attrs,
277                                sig: Some(item.kind),
278                                node_id: item.node_id,
279                            },
280                        )
281                    },
282                )?;
283            }
284            surface::DetachedItemKind::TraitImpl(trait_impl) => {
285                self.collect_trait_impl(owner_id, item.node_id, item.attrs, trait_impl)?;
286            }
287        };
288        Ok(())
289    }
290
291    fn collect_trait(
292        &mut self,
293        owner_id: OwnerId,
294        node_id: NodeId,
295        attrs: Vec<surface::Attr>,
296        trait_def: surface::DetachedTrait,
297    ) -> Result {
298        // 1. Collect the associated-refinements
299        self.inner.insert_item(
300            owner_id,
301            surface::Item {
302                attrs,
303                kind: surface::ItemKind::Trait(surface::Trait {
304                    generics: None,
305                    assoc_refinements: trait_def.refts,
306                }),
307                node_id,
308            },
309        )?;
310
311        // 2. Collect the method specifications
312        let tcx = self.inner.tcx;
313        let assoc_items = tcx.associated_items(owner_id.def_id).in_definition_order();
314        self.collect_assoc_methods(trait_def.items, assoc_items, |this, owner_id, item| {
315            this.inner.insert_trait_item(
316                owner_id,
317                surface::TraitItemFn {
318                    attrs: item.attrs,
319                    sig: Some(item.kind),
320                    node_id: item.node_id,
321                },
322            )
323        })
324    }
325
326    fn collect_trait_impl(
327        &mut self,
328        owner_id: OwnerId,
329        node_id: NodeId,
330        attrs: Vec<surface::Attr>,
331        trait_impl: surface::DetachedTraitImpl,
332    ) -> Result {
333        // 1. Collect the associated-refinements
334        self.inner.insert_item(
335            owner_id,
336            surface::Item {
337                attrs,
338                kind: surface::ItemKind::Impl(surface::Impl {
339                    generics: None,
340                    assoc_refinements: trait_impl.refts,
341                }),
342                node_id,
343            },
344        )?;
345
346        // 2. Collect the method specifications
347        let tcx = self.inner.tcx;
348        let assoc_items = tcx.associated_items(owner_id.def_id).in_definition_order();
349        self.collect_assoc_methods(trait_impl.items, assoc_items, |this, owner_id, item| {
350            this.inner.insert_impl_item(
351                owner_id,
352                surface::ImplItemFn {
353                    attrs: item.attrs,
354                    sig: Some(item.kind),
355                    node_id: item.node_id,
356                },
357            )
358        })
359    }
360
361    fn collect_assoc_methods(
362        &mut self,
363        methods: Vec<DetachedItem<surface::FnSig>>,
364        assoc_items: impl Iterator<Item = &'tcx AssocItem>,
365        mut insert_item: impl FnMut(&mut Self, OwnerId, DetachedItem<surface::FnSig>) -> Result,
366    ) -> Result {
367        let mut table: HashMap<Symbol, DetachedItem<(surface::FnSig, Option<DefId>)>> =
368            HashMap::default();
369        // 1. make a table of the impl-items
370        for item in methods {
371            let name = path_to_symbol(&item.path);
372            let span = item.path.span;
373            if let Entry::Occupied(_) = table.entry(name) {
374                return Err(self
375                    .inner
376                    .errors
377                    .emit(errors::MultipleSpecifications { name, span }));
378            } else {
379                table.insert(name, item.map_kind(|spec| (spec, None)));
380            }
381        }
382        // 2. walk over all the assoc-items to resolve names
383        for item in assoc_items {
384            if let AssocKind::Fn { name, .. } = item.kind
385                && let Some(val) = table.get_mut(&name)
386                && val.kind.1.is_none()
387            {
388                val.kind.1 = Some(item.def_id);
389            }
390        }
391        // 3. Attach the `fn_sig` to the resolved `DefId`
392        for (_name, item) in table {
393            let Some(def_id) = item.kind.1 else {
394                return Err(self
395                    .inner
396                    .errors
397                    .emit(errors::UnresolvedSpecification::new(&item.path, "identifier")));
398            };
399            if let Some(def_id) = self.unwrap_def_id(&def_id)? {
400                dbg::hyperlink!(self.inner.tcx, item.path.span, self.inner.tcx.def_span(def_id));
401                let owner_id = self.inner.tcx.local_def_id_to_hir_id(def_id).owner;
402                insert_item(self, owner_id, item.map_kind(|k| k.0))?;
403            }
404        }
405        Ok(())
406    }
407}
flux_driver/collector/detached_specs.rs

flux_driver/collector/
detached_specs.rs
