1use flux_common::index::IndexGen;
2use flux_errors::Errors;
3use flux_middle::fhir::{
4 self,
5 Namespace::{ReftNS, TypeNS, ValueNS},
6 PartialRes, Res,
7};
8use flux_syntax::{
9 surface::{self, Ident, NodeId, visit::Visitor as _},
10 walk_list,
11};
12use rustc_data_structures::unord::{ExtendUnord, UnordMap, UnordSet};
13use rustc_span::{ErrorGuaranteed, Span};
14
15use super::{BindingSource, CrateResolver, ResolveError, RibKind, Segment, errors::UnresolvedName};
16
17type Result<T = ()> = std::result::Result<T, ErrorGuaranteed>;
18
19pub(crate) enum Scope<'a> {
22 Implicit { kind: BinderKind, node_id: NodeId, params: &'a [surface::RefineParam] },
25 Explicit(&'a [surface::RefineParam]),
28 Opaque,
31}
32
33#[derive(Clone, Copy, PartialEq, Eq, Debug)]
34pub(crate) enum BinderKind {
35 FnInput,
39 FnOutput,
41 Variant,
43}
44
45impl BinderKind {
46 fn allows(self, bind_kind: surface::BindKind) -> bool {
48 match bind_kind {
49 surface::BindKind::At => {
50 matches!(self, Self::FnInput | Self::Variant)
51 }
52 surface::BindKind::Pound => self == Self::FnOutput,
53 }
54 }
55
56 fn rib_kind(self) -> RibKind {
57 match self {
58 BinderKind::Variant => RibKind::Variant,
59 BinderKind::FnInput | BinderKind::FnOutput => RibKind::Misc,
60 }
61 }
62}
63
64pub(crate) trait ScopedVisitor<'genv, 'tcx: 'genv>: Sized {
65 fn resolver(&self) -> &CrateResolver<'genv, 'tcx>;
66
67 fn enter_scope(&mut self, scope: &Scope);
68 fn exit_scope(&mut self, scope: &Scope);
69
70 fn wrap(self) -> ScopedVisitorWrapper<Self> {
71 ScopedVisitorWrapper(self)
72 }
73
74 fn on_implicit_param(&mut self, _ident: Ident, _kind: fhir::ParamKind, _node_id: NodeId) {}
75 fn on_path<S: Segment>(&mut self, _segments: &[S], _node_id: NodeId, _span: Span) {}
76 fn on_sort_path(&mut self, _path: &surface::SortPath) {}
77}
78
79pub(crate) struct ScopedVisitorWrapper<V>(V);
80
81impl<'genv, 'tcx: 'genv, V: ScopedVisitor<'genv, 'tcx>> ScopedVisitorWrapper<V> {
82 fn with_scope(&mut self, scope: Scope<'_>, f: impl FnOnce(&mut Self)) {
83 self.0.enter_scope(&scope);
84 f(self);
85 self.0.exit_scope(&scope);
86 }
87
88 fn is_box(&self, segment: &surface::PathSegment) -> bool {
89 self.resolver()
90 .output
91 .path_res_map
92 .get(&segment.node_id)
93 .map(|r| r.is_box(self.resolver().genv.tcx()))
94 .unwrap_or(false)
95 }
96
97 fn visit_let_decls<'v>(&mut self, decls: &'v [surface::LetDecl], body: &'v surface::Expr) {
100 let Some((decl, rest)) = decls.split_first() else {
101 self.visit_expr(body);
102 return;
103 };
104 self.visit_expr(&decl.init);
105 self.with_scope(Scope::Explicit(std::slice::from_ref(&decl.param)), |this| {
106 this.visit_refine_param(&decl.param);
107 this.visit_let_decls(rest, body);
108 });
109 }
110}
111
112impl<V> std::ops::Deref for ScopedVisitorWrapper<V> {
113 type Target = V;
114
115 fn deref(&self) -> &Self::Target {
116 &self.0
117 }
118}
119impl<V> std::ops::DerefMut for ScopedVisitorWrapper<V> {
120 fn deref_mut(&mut self) -> &mut Self::Target {
121 &mut self.0
122 }
123}
124
125impl<'v, 'genv, 'tcx: 'genv, V: ScopedVisitor<'genv, 'tcx>> surface::visit::Visitor<'v>
126 for ScopedVisitorWrapper<V>
127{
128 fn visit_trait_assoc_reft(&mut self, assoc_reft: &'v surface::TraitAssocReft) {
129 self.with_scope(Scope::Explicit(&assoc_reft.params), |this| {
130 surface::visit::walk_trait_assoc_reft(this, assoc_reft);
131 });
132 }
133
134 fn visit_impl_assoc_reft(&mut self, assoc_reft: &'v surface::ImplAssocReft) {
135 self.with_scope(Scope::Explicit(&assoc_reft.params), |this| {
136 surface::visit::walk_impl_assoc_reft(this, assoc_reft);
137 });
138 }
139
140 fn visit_qualifier(&mut self, qualifier: &'v surface::Qualifier) {
141 self.with_scope(Scope::Explicit(&qualifier.params), |this| {
142 surface::visit::walk_qualifier(this, qualifier);
143 });
144 }
145
146 fn visit_defn(&mut self, defn: &'v surface::SpecFunc) {
147 self.with_scope(Scope::Explicit(&defn.params), |this| {
148 surface::visit::walk_defn(this, defn);
149 });
150 }
151
152 fn visit_primop_prop(&mut self, prop: &'v surface::PrimOpProp) {
153 self.with_scope(Scope::Explicit(&prop.params), |this| {
154 surface::visit::walk_primop_prop(this, prop);
155 });
156 }
157
158 fn visit_axiom(&mut self, axiom: &'v surface::Axiom) {
159 self.with_scope(Scope::Explicit(&axiom.params), |this| {
160 surface::visit::walk_axiom(this, axiom);
161 });
162 }
163
164 fn visit_ty_alias(&mut self, ty_alias: &'v surface::TyAlias) {
165 self.visit_generics(&ty_alias.generics);
167 self.with_scope(Scope::Explicit(&ty_alias.params), |this| {
168 walk_list!(this, visit_refine_param, &ty_alias.params);
169 this.with_scope(Scope::Explicit(ty_alias.index.as_slice()), |this| {
170 walk_list!(this, visit_refine_param, &ty_alias.index);
171 this.visit_ty(&ty_alias.ty);
172 });
173 });
174 }
175
176 fn visit_struct_def(&mut self, struct_def: &'v surface::StructDef) {
177 if let Some(generics) = &struct_def.generics {
179 self.visit_generics(generics);
180 }
181 let refined_by = struct_def.refined_by.as_deref().unwrap_or_default();
182 self.with_scope(Scope::Explicit(refined_by), |this| {
183 walk_list!(this, visit_refine_param, refined_by);
184 walk_list!(this, visit_expr, &struct_def.invariants);
185 struct_def.fields.iter().flatten().for_each(|field| {
186 this.visit_ty(field);
187 });
188 });
189 }
190
191 fn visit_enum_def(&mut self, enum_def: &'v surface::EnumDef) {
192 let refined_by = enum_def.refined_by.as_deref().unwrap_or_default();
193 self.with_scope(Scope::Explicit(refined_by), |this| {
194 surface::visit::walk_enum_def(this, enum_def);
195 });
196 }
197
198 fn visit_variant(&mut self, variant: &'v surface::VariantDef) {
199 let scope =
200 Scope::Implicit { kind: BinderKind::Variant, node_id: variant.node_id, params: &[] };
201 self.with_scope(scope, |this| {
202 surface::visit::walk_variant(this, variant);
203 });
204 }
205
206 fn visit_trait_ref(&mut self, trait_ref: &'v surface::TraitRef) {
207 match trait_ref.as_fn_trait_ref() {
208 Some((in_arg, out_arg)) => {
209 let scope = Scope::Implicit {
210 kind: BinderKind::FnInput,
211 node_id: trait_ref.node_id,
212 params: &[],
213 };
214 self.with_scope(scope, |this| {
215 surface::visit::walk_generic_arg(this, in_arg);
216 this.with_scope(Scope::Opaque, |this| {
217 surface::visit::walk_generic_arg(this, out_arg);
218 });
219 });
220 }
221 None => {
222 self.with_scope(Scope::Opaque, |this| {
223 surface::visit::walk_trait_ref(this, trait_ref);
224 });
225 }
226 }
227 }
228
229 fn visit_variant_ret(&mut self, ret: &'v surface::VariantRet) {
230 self.with_scope(Scope::Opaque, |this| {
231 surface::visit::walk_variant_ret(this, ret);
232 });
233 }
234
235 fn visit_generics(&mut self, generics: &'v surface::Generics) {
236 self.with_scope(Scope::Opaque, |this| {
237 surface::visit::walk_generics(this, generics);
238 });
239 }
240
241 fn visit_fn_sig(&mut self, fn_sig: &'v surface::FnSig) {
242 let scope = Scope::Implicit {
243 kind: BinderKind::FnInput,
244 node_id: fn_sig.node_id,
245 params: &fn_sig.params,
246 };
247 self.with_scope(scope, |this| {
248 surface::visit::walk_fn_sig(this, fn_sig);
249 });
250 }
251
252 fn visit_requires(&mut self, requires: &'v surface::Requires) {
253 self.with_scope(Scope::Explicit(&requires.params), |this| {
254 surface::visit::walk_requires(this, requires);
255 });
256 }
257
258 fn visit_fn_output(&mut self, output: &'v surface::FnOutput) {
259 let scope =
260 Scope::Implicit { kind: BinderKind::FnOutput, node_id: output.node_id, params: &[] };
261 self.with_scope(scope, |this| {
262 surface::visit::walk_fn_output(this, output);
263 });
264 }
265
266 fn visit_fn_input(&mut self, arg: &'v surface::FnInput) {
267 match arg {
268 surface::FnInput::Constr(bind, _, _, node_id) => {
269 self.on_implicit_param(*bind, fhir::ParamKind::Colon, *node_id);
270 }
271 surface::FnInput::StrgRef(loc, _, node_id) => {
272 self.on_implicit_param(*loc, fhir::ParamKind::Loc, *node_id);
273 }
274 surface::FnInput::Ty(bind, ty, node_id) => {
275 if let &Some(bind) = bind {
276 let param_kind = if let surface::TyKind::Base(_) = &ty.kind {
277 fhir::ParamKind::Colon
278 } else {
279 fhir::ParamKind::Error
280 };
281 self.on_implicit_param(bind, param_kind, *node_id);
282 }
283 }
284 }
285 surface::visit::walk_fn_input(self, arg);
286 }
287
288 fn visit_ensures(&mut self, constraint: &'v surface::Ensures) {
289 if let surface::Ensures::Type(loc, _, node_id) = constraint {
290 self.on_path(std::slice::from_ref(loc), *node_id, loc.span);
291 }
292 surface::visit::walk_ensures(self, constraint);
293 }
294
295 fn visit_refine_arg(&mut self, arg: &'v surface::RefineArg) {
296 match arg {
297 surface::RefineArg::Bind(ident, kind, _, node_id) => {
298 let kind = match kind {
299 surface::BindKind::At => fhir::ParamKind::At,
300 surface::BindKind::Pound => fhir::ParamKind::Pound,
301 };
302 self.on_implicit_param(*ident, kind, *node_id);
303 }
304 surface::RefineArg::Abs(params, ..) => {
305 self.with_scope(Scope::Explicit(params), |this| {
306 surface::visit::walk_refine_arg(this, arg);
307 });
308 }
309 surface::RefineArg::Expr(expr) => self.visit_expr(expr),
310 }
311 }
312
313 fn visit_path(&mut self, path: &'v surface::Path) {
314 for arg in &path.refine {
315 self.with_scope(Scope::Opaque, |this| this.visit_refine_arg(arg));
316 }
317 walk_list!(self, visit_path_segment, &path.segments);
318 }
319
320 fn visit_path_segment(&mut self, segment: &'v surface::PathSegment) {
321 let is_box = self.is_box(segment);
322 for (i, arg) in segment.args.iter().enumerate() {
323 if is_box && i == 0 {
324 self.visit_generic_arg(arg);
325 } else {
326 self.with_scope(Scope::Opaque, |this| this.visit_generic_arg(arg));
327 }
328 }
329 }
330
331 fn visit_ty(&mut self, ty: &'v surface::Ty) {
332 let node_id = ty.node_id;
333 match &ty.kind {
334 surface::TyKind::Exists { bind, .. } => {
335 let param = surface::RefineParam {
336 ident: *bind,
337 mode: None,
338 sort: surface::Sort::Infer,
339 node_id,
340 span: bind.span,
341 };
342 self.with_scope(Scope::Explicit(std::slice::from_ref(¶m)), |this| {
343 surface::visit::walk_ty(this, ty);
344 });
345 }
346 surface::TyKind::GeneralExists { params, .. } => {
347 self.with_scope(Scope::Explicit(params), |this| {
348 surface::visit::walk_ty(this, ty);
349 });
350 }
351 surface::TyKind::BareFn(bare_fn) => {
352 let scope = Scope::Implicit {
353 kind: BinderKind::FnInput,
354 node_id: bare_fn.node_id,
355 params: &[],
356 };
357 self.with_scope(scope, |this| {
358 walk_list!(this, visit_ty, &bare_fn.inputs);
359 this.with_scope(Scope::Opaque, |this| {
360 this.visit_fn_ret_ty(&bare_fn.output);
361 });
362 });
363 }
364 surface::TyKind::Array(..) => {
365 self.with_scope(Scope::Opaque, |this| {
366 surface::visit::walk_ty(this, ty);
367 });
368 }
369 _ => surface::visit::walk_ty(self, ty),
370 }
371 }
372
373 fn visit_bty(&mut self, bty: &'v surface::BaseTy) {
374 match &bty.kind {
375 surface::BaseTyKind::Slice(_) | surface::BaseTyKind::Ptr(..) => {
376 self.with_scope(Scope::Opaque, |this| {
377 surface::visit::walk_bty(this, bty);
378 });
379 }
380 surface::BaseTyKind::Path(..) => {
381 surface::visit::walk_bty(self, bty);
382 }
383 }
384 }
385
386 fn visit_expr(&mut self, expr: &'v surface::Expr) {
387 match &expr.kind {
388 surface::ExprKind::Quant(_, param, _, body) => {
389 self.with_scope(Scope::Explicit(std::slice::from_ref(param)), |this| {
390 this.visit_refine_param(param);
391 this.visit_expr(body);
392 });
393 }
394 surface::ExprKind::Block(decls, body) => self.visit_let_decls(decls, body),
395 _ => surface::visit::walk_expr(self, expr),
396 }
397 }
398
399 fn visit_path_expr(&mut self, path: &'v surface::ExprPath) {
400 self.on_path(&path.segments, path.node_id, path.span);
401 }
402
403 fn visit_sort_path(&mut self, path: &'v surface::SortPath) {
404 self.on_sort_path(path);
405 surface::visit::walk_sort_path(self, path);
406 }
407}
408
409#[derive(Clone, Copy)]
410pub(crate) struct ParamDef {
411 ident: Ident,
412 kind: fhir::ParamKind,
413 param_id: NodeId,
414}
415
416pub(crate) type ImplicitParams = UnordMap<NodeId, Vec<ParamDef>>;
417
418pub(crate) struct ImplicitParamCollector<'a, 'genv, 'tcx> {
421 resolver: &'a CrateResolver<'genv, 'tcx>,
422 scopes: Vec<Option<(BinderKind, NodeId)>>,
425 implicit_params: ImplicitParams,
426 errors: Errors<'genv>,
427}
428
429impl<'a, 'genv, 'tcx> ImplicitParamCollector<'a, 'genv, 'tcx> {
430 pub(crate) fn run(
431 resolver: &'a CrateResolver<'genv, 'tcx>,
432 f: impl FnOnce(&mut ScopedVisitorWrapper<Self>),
433 ) -> Result<ImplicitParams> {
434 let errors = Errors::new(resolver.genv.sess());
435 let mut vis =
436 Self { resolver, scopes: vec![], implicit_params: Default::default(), errors }.wrap();
437 f(&mut vis);
438 vis.0.errors.to_result()?;
439 Ok(vis.0.implicit_params)
440 }
441}
442
443impl<'genv, 'tcx> ScopedVisitor<'genv, 'tcx> for ImplicitParamCollector<'_, 'genv, 'tcx> {
444 fn resolver(&self) -> &CrateResolver<'genv, 'tcx> {
445 self.resolver
446 }
447
448 fn enter_scope(&mut self, scope: &Scope) {
449 let binder_scope = match scope {
450 Scope::Implicit { kind, node_id, .. } => Some((*kind, *node_id)),
451 Scope::Explicit(_) | Scope::Opaque => None,
452 };
453 self.scopes.push(binder_scope);
454 }
455
456 fn exit_scope(&mut self, _: &Scope) {
457 self.scopes.pop();
458 }
459
460 fn on_implicit_param(&mut self, ident: Ident, kind: fhir::ParamKind, param_id: NodeId) {
461 let Some(&binder_scope) = self.scopes.last() else { return };
462 let bind_kind = match kind {
463 fhir::ParamKind::At => Some(surface::BindKind::At),
464 fhir::ParamKind::Pound => Some(surface::BindKind::Pound),
465 _ => None,
466 };
467 if let Some(bind_kind) = bind_kind
468 && !binder_scope.is_some_and(|(kind, _)| kind.allows(bind_kind))
469 {
470 self.errors
471 .emit(errors::IllegalBinder::new(ident.span, bind_kind));
472 return;
473 }
474 if let Some((_, node_id)) = binder_scope {
475 self.implicit_params
476 .entry(node_id)
477 .or_default()
478 .push(ParamDef { ident, kind, param_id });
479 }
480 }
481}
482
483pub(crate) struct RefinementResolver<'a, 'genv, 'tcx> {
484 param_defs: Vec<ParamDef>,
485 used_params: UnordSet<NodeId>,
488 implicit_params: ImplicitParams,
490 resolver: &'a mut CrateResolver<'genv, 'tcx>,
491 errors: Errors<'genv>,
492}
493
494impl<'a, 'genv, 'tcx> RefinementResolver<'a, 'genv, 'tcx> {
495 pub(crate) fn run(
496 resolver: &'a mut CrateResolver<'genv, 'tcx>,
497 implicit_params: ImplicitParams,
498 f: impl FnOnce(&mut ScopedVisitorWrapper<Self>),
499 ) -> Result {
500 let errors = Errors::new(resolver.genv.sess());
501 let mut wrapper = Self {
502 resolver,
503 param_defs: Default::default(),
504 used_params: Default::default(),
505 implicit_params,
506 errors,
507 }
508 .wrap();
509 f(&mut wrapper);
510 wrapper.0.finish()
511 }
512
513 fn define_param(&mut self, def: ParamDef) {
514 self.param_defs.push(def);
515 self.resolver.define_res_in(
516 Res::Param(def.kind, def.param_id),
517 ReftNS,
518 BindingSource::Explicit(def.ident),
519 );
520 }
521
522 fn define_implicit_params(&mut self, node_id: NodeId) {
523 for def in self
524 .implicit_params
525 .get(&node_id)
526 .cloned()
527 .unwrap_or_default()
528 {
529 self.define_param(def);
530 }
531 }
532
533 fn define_explicit_params(&mut self, params: &[surface::RefineParam]) {
534 for param in params {
535 self.define_param(ParamDef {
536 ident: param.ident,
537 kind: fhir::ParamKind::Explicit(param.mode.map(Into::into)),
538 param_id: param.node_id,
539 });
540 }
541 }
542
543 fn resolve_path<S: Segment>(&mut self, segments: &[S], node_id: NodeId, span: Span) {
544 match self.try_resolve_expr_with_ribs(segments) {
545 Ok(res) => self.record_expr_res(node_id, res, segments.last().unwrap().ident()),
546 Err(err) => {
547 err.emit(&self.errors, || UnresolvedName::new(span, segments, "value"));
548 }
549 }
550 }
551
552 fn record_expr_res(&mut self, node_id: NodeId, res: PartialRes<NodeId>, ident: Ident) {
553 if let Some(Res::Param(fhir::ParamKind::Error, _)) = res.full_res() {
555 self.errors.emit(errors::InvalidUnrefinedParam::new(ident));
556 }
557 if let Res::Param(_, param_id) = res.base_res() {
558 self.used_params.insert(param_id);
559 }
560 self.resolver.output.path_res_map.insert(node_id, res);
561 }
562
563 fn try_resolve_expr_with_ribs<S: Segment>(
564 &mut self,
565 segments: &[S],
566 ) -> std::result::Result<PartialRes<NodeId>, ResolveError> {
567 let mut ambiguity = None;
576 for ns in [ReftNS, ValueNS, TypeNS] {
577 match self.resolver.resolve_path_with_ribs(segments, ns) {
578 Ok(partial_res) => return Ok(partial_res),
579 Err(ResolveError::Ambiguous(amb)) => ambiguity = ambiguity.or(Some(amb)),
580 Err(ResolveError::NotFound) => {}
581 Err(err @ ResolveError::OuterGeneric(_)) => return Err(err),
582 }
583 }
584 Err(ambiguity.map_or(ResolveError::NotFound, ResolveError::Ambiguous))
585 }
586
587 fn resolve_sort_path(&mut self, path: &surface::SortPath) {
588 let res = match self.resolver.resolve_path_with_ribs(&path.segments, TypeNS) {
593 Ok(res) => res,
594 Err(err) => {
595 err.emit(&self.errors, || UnresolvedName::new(path.span(), &path.segments, "sort"));
596 PartialRes::new(fhir::Res::Err)
597 }
598 };
599 self.resolver.output.path_res_map.insert(path.node_id, res);
600 }
601
602 fn finish(self) -> Result {
603 let param_id_gen = IndexGen::new();
604 let output = &mut self.resolver.output;
605 for def in self.param_defs {
606 match def.kind {
607 fhir::ParamKind::Colon if !self.used_params.contains(&def.param_id) => continue,
608 fhir::ParamKind::Error => continue,
609 _ => {}
610 }
611 output
612 .param_res_map
613 .insert(def.param_id, (param_id_gen.fresh(), def.kind));
614 }
615 output
616 .implicit_params
617 .extend_unord(self.implicit_params.into_items().map(|(scope, params)| {
618 let params = params
619 .into_iter()
620 .filter(|def| output.param_res_map.contains_key(&def.param_id))
621 .map(|def| (def.ident, def.param_id))
622 .collect();
623 (scope, params)
624 }));
625 self.errors.to_result()
626 }
627}
628
629impl<'genv, 'tcx> ScopedVisitor<'genv, 'tcx> for RefinementResolver<'_, 'genv, 'tcx> {
630 fn resolver(&self) -> &CrateResolver<'genv, 'tcx> {
631 self.resolver
632 }
633
634 fn enter_scope(&mut self, scope: &Scope) {
635 match *scope {
636 Scope::Implicit { kind, node_id, params } => {
637 self.resolver.push_rib(ReftNS, kind.rib_kind());
638 self.define_explicit_params(params);
639 self.define_implicit_params(node_id);
640 }
641 Scope::Explicit(params) => {
642 self.resolver.push_rib(ReftNS, RibKind::Misc);
643 self.define_explicit_params(params);
644 }
645 Scope::Opaque => {}
647 }
648 }
649
650 fn exit_scope(&mut self, scope: &Scope) {
651 match scope {
652 Scope::Implicit { .. } | Scope::Explicit(_) => self.resolver.pop_rib(ReftNS),
653 Scope::Opaque => {}
654 }
655 }
656
657 fn on_path<S: Segment>(&mut self, segments: &[S], node_id: NodeId, span: Span) {
658 self.resolve_path(segments, node_id, span);
659 }
660
661 fn on_sort_path(&mut self, path: &surface::SortPath) {
662 self.resolve_sort_path(path);
663 }
664}
665
666mod errors {
667 use flux_errors::E0999;
668 use flux_macros::Diagnostic;
669 use flux_syntax::surface;
670 use rustc_span::{Span, symbol::Ident};
671
672 #[derive(Diagnostic)]
673 #[diag("invalid use of refinement parameter", code = E0999)]
674 pub(super) struct InvalidUnrefinedParam {
675 #[primary_span]
676 #[label("parameter `{$var}` refers to a type with no indices")]
677 span: Span,
678 var: Ident,
679 }
680
681 impl InvalidUnrefinedParam {
682 pub(super) fn new(var: Ident) -> Self {
683 Self { var, span: var.span }
684 }
685 }
686
687 #[derive(Diagnostic)]
688 #[diag("illegal binder", code = E0999)]
689 pub(super) struct IllegalBinder {
690 #[primary_span]
691 #[label("`{$kind}` binder not allowed in this position")]
692 span: Span,
693 kind: &'static str,
694 }
695
696 impl IllegalBinder {
697 pub(super) fn new(span: Span, kind: surface::BindKind) -> Self {
698 Self { span, kind: kind.token_str() }
699 }
700 }
701}