flux_driver/collector/
detached_specs.rs

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

flux_driver/collector/
detached_specs.rs
