flux_middle/rty/
fold.rs

1//! This modules follows the implementation of folding in rustc. For more information read the
2//! documentation in `rustc_type_ir::fold`.
3
4use std::ops::ControlFlow;
5
6use flux_arc_interner::{Internable, List};
7use flux_common::bug;
8use itertools::Itertools;
9use rustc_data_structures::fx::FxHashMap;
10use rustc_hash::FxHashSet;
11use rustc_type_ir::{BoundVar, DebruijnIndex, INNERMOST};
12
13use super::{
14    BaseTy, Binder, BoundVariableKinds, Const, EVid, EarlyReftParam, Ensures, Expr, ExprKind,
15    GenericArg, Name, OutlivesPredicate, PolyFuncSort, PtrKind, ReBound, ReErased, Region, Sort,
16    SubsetTy, Ty, TyKind, TyOrBase, normalize::Normalizer,
17};
18use crate::{
19    global_env::GlobalEnv,
20    rty::{BoundReft, BoundRegion, Var, VariantSig, expr::HoleKind},
21};
22
23pub trait TypeVisitor: Sized {
24    type BreakTy = !;
25
26    fn enter_binder(&mut self, _vars: &BoundVariableKinds) {}
27
28    fn exit_binder(&mut self) {}
29
30    fn visit_expr(&mut self, expr: &Expr) -> ControlFlow<Self::BreakTy> {
31        expr.super_visit_with(self)
32    }
33
34    fn visit_sort(&mut self, sort: &Sort) -> ControlFlow<Self::BreakTy> {
35        sort.super_visit_with(self)
36    }
37
38    fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::BreakTy> {
39        ty.super_visit_with(self)
40    }
41
42    fn visit_bty(&mut self, bty: &BaseTy) -> ControlFlow<Self::BreakTy> {
43        bty.super_visit_with(self)
44    }
45}
46
47pub trait FallibleTypeFolder: Sized {
48    type Error;
49
50    fn try_enter_binder(&mut self, _vars: &BoundVariableKinds) {}
51
52    fn try_exit_binder(&mut self) {}
53
54    fn try_fold_sort(&mut self, sort: &Sort) -> Result<Sort, Self::Error> {
55        sort.try_super_fold_with(self)
56    }
57
58    fn try_fold_ty(&mut self, ty: &Ty) -> Result<Ty, Self::Error> {
59        ty.try_super_fold_with(self)
60    }
61
62    fn try_fold_bty(&mut self, bty: &BaseTy) -> Result<BaseTy, Self::Error> {
63        bty.try_super_fold_with(self)
64    }
65
66    fn try_fold_subset_ty(&mut self, constr: &SubsetTy) -> Result<SubsetTy, Self::Error> {
67        constr.try_super_fold_with(self)
68    }
69
70    fn try_fold_region(&mut self, re: &Region) -> Result<Region, Self::Error> {
71        Ok(*re)
72    }
73
74    fn try_fold_const(&mut self, c: &Const) -> Result<Const, Self::Error> {
75        c.try_super_fold_with(self)
76    }
77
78    fn try_fold_expr(&mut self, expr: &Expr) -> Result<Expr, Self::Error> {
79        expr.try_super_fold_with(self)
80    }
81}
82
83pub trait TypeFolder: FallibleTypeFolder<Error = !> {
84    fn enter_binder(&mut self, _vars: &BoundVariableKinds) {}
85
86    fn exit_binder(&mut self) {}
87
88    fn fold_sort(&mut self, sort: &Sort) -> Sort {
89        sort.super_fold_with(self)
90    }
91
92    fn fold_ty(&mut self, ty: &Ty) -> Ty {
93        ty.super_fold_with(self)
94    }
95
96    fn fold_bty(&mut self, bty: &BaseTy) -> BaseTy {
97        bty.super_fold_with(self)
98    }
99
100    fn fold_subset_ty(&mut self, constr: &SubsetTy) -> SubsetTy {
101        constr.super_fold_with(self)
102    }
103
104    fn fold_region(&mut self, re: &Region) -> Region {
105        *re
106    }
107
108    fn fold_const(&mut self, c: &Const) -> Const {
109        c.super_fold_with(self)
110    }
111
112    fn fold_expr(&mut self, expr: &Expr) -> Expr {
113        expr.super_fold_with(self)
114    }
115}
116
117impl<F> FallibleTypeFolder for F
118where
119    F: TypeFolder,
120{
121    type Error = !;
122
123    fn try_enter_binder(&mut self, vars: &BoundVariableKinds) {
124        self.enter_binder(vars);
125    }
126
127    fn try_exit_binder(&mut self) {
128        self.exit_binder();
129    }
130
131    fn try_fold_sort(&mut self, sort: &Sort) -> Result<Sort, Self::Error> {
132        Ok(self.fold_sort(sort))
133    }
134
135    fn try_fold_ty(&mut self, ty: &Ty) -> Result<Ty, Self::Error> {
136        Ok(self.fold_ty(ty))
137    }
138
139    fn try_fold_bty(&mut self, bty: &BaseTy) -> Result<BaseTy, Self::Error> {
140        Ok(self.fold_bty(bty))
141    }
142
143    fn try_fold_subset_ty(&mut self, ty: &SubsetTy) -> Result<SubsetTy, Self::Error> {
144        Ok(self.fold_subset_ty(ty))
145    }
146
147    fn try_fold_region(&mut self, re: &Region) -> Result<Region, Self::Error> {
148        Ok(self.fold_region(re))
149    }
150
151    fn try_fold_const(&mut self, c: &Const) -> Result<Const, Self::Error> {
152        Ok(self.fold_const(c))
153    }
154
155    fn try_fold_expr(&mut self, expr: &Expr) -> Result<Expr, Self::Error> {
156        Ok(self.fold_expr(expr))
157    }
158}
159
160pub trait TypeVisitable: Sized {
161    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy>;
162
163    fn has_escaping_bvars(&self) -> bool {
164        self.has_escaping_bvars_at_or_above(INNERMOST)
165    }
166
167    /// Returns `true` if `self` has any late-bound vars that are either
168    /// bound by `binder` or bound by some binder outside of `binder`.
169    /// If `binder` is `ty::INNERMOST`, this indicates whether
170    /// there are any late-bound vars that appear free.
171    fn has_escaping_bvars_at_or_above(&self, binder: DebruijnIndex) -> bool {
172        struct HasEscapingVars {
173            /// Anything bound by `outer_index` or "above" is escaping.
174            outer_index: DebruijnIndex,
175        }
176
177        impl TypeVisitor for HasEscapingVars {
178            type BreakTy = ();
179
180            fn enter_binder(&mut self, _: &BoundVariableKinds) {
181                self.outer_index.shift_in(1);
182            }
183
184            fn exit_binder(&mut self) {
185                self.outer_index.shift_out(1);
186            }
187
188            // TODO(nilehmann) keep track of the outermost binder to optimize this, i.e.,
189            // what rustc calls outer_exclusive_binder.
190            fn visit_expr(&mut self, expr: &Expr) -> ControlFlow<()> {
191                if let ExprKind::Var(Var::Bound(debruijn, _)) = expr.kind() {
192                    if *debruijn >= self.outer_index {
193                        ControlFlow::Break(())
194                    } else {
195                        ControlFlow::Continue(())
196                    }
197                } else {
198                    expr.super_visit_with(self)
199                }
200            }
201        }
202        let mut visitor = HasEscapingVars { outer_index: binder };
203        self.visit_with(&mut visitor).is_break()
204    }
205
206    /// Returns the set of all free variables.
207    /// For example, `Vec<i32[n]>{v : v > m}` returns `{n, m}`.
208    fn fvars(&self) -> FxHashSet<Name> {
209        struct CollectFreeVars(FxHashSet<Name>);
210
211        impl TypeVisitor for CollectFreeVars {
212            fn visit_expr(&mut self, e: &Expr) -> ControlFlow<Self::BreakTy> {
213                if let ExprKind::Var(Var::Free(name)) = e.kind() {
214                    self.0.insert(*name);
215                }
216                e.super_visit_with(self)
217            }
218        }
219
220        let mut collector = CollectFreeVars(FxHashSet::default());
221        let _ = self.visit_with(&mut collector);
222        collector.0
223    }
224
225    fn early_params(&self) -> FxHashSet<EarlyReftParam> {
226        struct CollectEarlyParams(FxHashSet<EarlyReftParam>);
227
228        impl TypeVisitor for CollectEarlyParams {
229            fn visit_expr(&mut self, e: &Expr) -> ControlFlow<Self::BreakTy> {
230                if let ExprKind::Var(Var::EarlyParam(param)) = e.kind() {
231                    self.0.insert(*param);
232                }
233                e.super_visit_with(self)
234            }
235        }
236
237        let mut collector = CollectEarlyParams(FxHashSet::default());
238        let _ = self.visit_with(&mut collector);
239        collector.0
240    }
241
242    /// Gives the indices of the provided bvars which:
243    ///   1. Only occur a single time.
244    ///   2. In their occurrence, are either
245    ///      a. The direct argument in an index (e.g. `exists b0. usize[b0]`)
246    ///      b. The direct argument of a constructor in an index (e.g.
247    ///      `exists b0. Vec<usize>[{len: b0}]`)
248    ///
249    /// This is to be used for "re-sugaring" existentials into surface syntax
250    /// that doesn't use existentials.
251    ///
252    /// For 2b., we do need to be careful to ensure that if a constructor has
253    /// multiple arguments, they _all_ are redundant bvars, e.g. as in
254    ///
255    ///     exists b0, b1. RMat<f32>[{rows: b0, cols: b1}]
256    ///
257    /// which may be rewritten as `RMat<f32>`,
258    /// versus the (unlikely) edge case
259    ///
260    ///     exists b0. RMat<f32>[{rows: b0, cols: b0}]
261    ///
262    /// for which the existential is now necessary.
263    ///
264    /// NOTE: this only applies to refinement bvars.
265    fn redundant_bvars(&self) -> FxHashSet<BoundVar> {
266        struct RedundantBVarFinder {
267            current_index: DebruijnIndex,
268            total_bvar_occurrences: FxHashMap<BoundVar, usize>,
269            bvars_appearing_in_index: FxHashSet<BoundVar>,
270        }
271
272        impl TypeVisitor for RedundantBVarFinder {
273            // Here we count all times we see a bvar
274            fn visit_expr(&mut self, e: &Expr) -> ControlFlow<Self::BreakTy> {
275                if let ExprKind::Var(Var::Bound(debruijn, BoundReft { var, .. })) = e.kind()
276                    && debruijn == &self.current_index
277                {
278                    self.total_bvar_occurrences
279                        .entry(*var)
280                        .and_modify(|count| {
281                            *count += 1;
282                        })
283                        .or_insert(1);
284                }
285                e.super_visit_with(self)
286            }
287
288            fn visit_ty(&mut self, ty: &Ty) -> ControlFlow<Self::BreakTy> {
289                // Here we check for bvars specifically as the direct arguments
290                // to an index or as the direct arguments to a Ctor in an index.
291                if let TyKind::Indexed(_bty, expr) = ty.kind() {
292                    match expr.kind() {
293                        ExprKind::Var(Var::Bound(debruijn, BoundReft { var, .. })) => {
294                            if debruijn == &self.current_index {
295                                self.bvars_appearing_in_index.insert(*var);
296                            }
297                        }
298                        ExprKind::Ctor(_ctor, exprs) => {
299                            exprs.iter().for_each(|expr| {
300                                if let ExprKind::Var(Var::Bound(debruijn, BoundReft { var, .. })) =
301                                    expr.kind()
302                                    && debruijn == &self.current_index
303                                {
304                                    self.bvars_appearing_in_index.insert(*var);
305                                }
306                            });
307                        }
308                        _ => {}
309                    }
310                }
311                ty.super_visit_with(self)
312            }
313
314            fn enter_binder(&mut self, _: &BoundVariableKinds) {
315                self.current_index.shift_in(1);
316            }
317
318            fn exit_binder(&mut self) {
319                self.current_index.shift_out(1);
320            }
321        }
322
323        let mut finder = RedundantBVarFinder {
324            current_index: INNERMOST,
325            total_bvar_occurrences: FxHashMap::default(),
326            bvars_appearing_in_index: FxHashSet::default(),
327        };
328        let _ = self.visit_with(&mut finder);
329
330        finder
331            .bvars_appearing_in_index
332            .into_iter()
333            .filter(|var_index| finder.total_bvar_occurrences.get(var_index) == Some(&1))
334            .collect()
335    }
336}
337
338pub trait TypeSuperVisitable: TypeVisitable {
339    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy>;
340}
341
342pub trait TypeFoldable: TypeVisitable {
343    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error>;
344
345    fn fold_with<F: TypeFolder>(&self, folder: &mut F) -> Self {
346        self.try_fold_with(folder).into_ok()
347    }
348
349    /// Normalize expressions by applying beta reductions for tuples and lambda abstractions.
350    fn normalize(&self, genv: GlobalEnv) -> Self {
351        self.fold_with(&mut Normalizer::new(genv, None))
352    }
353
354    /// Replaces all [holes] with the result of calling a closure. The closure takes a list with
355    /// all the *layers* of [bound] variables at the point the hole was found. Each layer corresponds
356    /// to the list of bound variables at that level. The list is ordered from outermost to innermost
357    /// binder, i.e., the last element is the binder closest to the hole.
358    ///
359    /// [holes]: ExprKind::Hole
360    /// [bound]: Binder
361    fn replace_holes(&self, f: impl FnMut(&[BoundVariableKinds], HoleKind) -> Expr) -> Self {
362        struct ReplaceHoles<F>(F, Vec<BoundVariableKinds>);
363
364        impl<F> TypeFolder for ReplaceHoles<F>
365        where
366            F: FnMut(&[BoundVariableKinds], HoleKind) -> Expr,
367        {
368            fn enter_binder(&mut self, vars: &BoundVariableKinds) {
369                self.1.push(vars.clone());
370            }
371
372            fn exit_binder(&mut self) {
373                self.1.pop();
374            }
375
376            fn fold_expr(&mut self, e: &Expr) -> Expr {
377                if let ExprKind::Hole(kind) = e.kind() {
378                    self.0(&self.1, kind.clone())
379                } else {
380                    e.super_fold_with(self)
381                }
382            }
383        }
384
385        self.fold_with(&mut ReplaceHoles(f, vec![]))
386    }
387
388    /// Remove all refinements and turn each underlying [`BaseTy`] into a [`TyKind::Exists`] with a
389    /// [`TyKind::Constr`] and a [`hole`]. For example, `Vec<{v. i32[v] | v > 0}>[n]` becomes
390    /// `{n. Vec<{v. i32[v] | *}>[n] | *}`.
391    ///
392    /// [`hole`]: ExprKind::Hole
393    fn with_holes(&self) -> Self {
394        struct WithHoles;
395
396        impl TypeFolder for WithHoles {
397            fn fold_ty(&mut self, ty: &Ty) -> Ty {
398                if let Some(bty) = ty.as_bty_skipping_existentials() {
399                    Ty::exists_with_constr(bty.fold_with(self), Expr::hole(HoleKind::Pred))
400                } else {
401                    ty.super_fold_with(self)
402                }
403            }
404
405            fn fold_subset_ty(&mut self, constr: &SubsetTy) -> SubsetTy {
406                SubsetTy::new(constr.bty.clone(), constr.idx.clone(), Expr::hole(HoleKind::Pred))
407            }
408        }
409
410        self.fold_with(&mut WithHoles)
411    }
412
413    fn replace_evars(&self, f: &mut impl FnMut(EVid) -> Option<Expr>) -> Result<Self, EVid> {
414        struct Folder<F>(F);
415        impl<F: FnMut(EVid) -> Option<Expr>> FallibleTypeFolder for Folder<F> {
416            type Error = EVid;
417
418            fn try_fold_expr(&mut self, expr: &Expr) -> Result<Expr, Self::Error> {
419                if let ExprKind::Var(Var::EVar(evid)) = expr.kind() {
420                    if let Some(sol) = (self.0)(*evid) { Ok(sol.clone()) } else { Err(*evid) }
421                } else {
422                    expr.try_super_fold_with(self)
423                }
424            }
425        }
426
427        self.try_fold_with(&mut Folder(f))
428    }
429
430    /// Shift the [`BoundVar`] of all the the variables at index [`INNERMOST`] "horizontally" to the
431    /// right by the specified `amount`.
432    fn shift_horizontally(&self, amount: usize) -> Self {
433        struct Shifter {
434            current_index: DebruijnIndex,
435            amount: usize,
436        }
437
438        impl TypeFolder for Shifter {
439            fn enter_binder(&mut self, _: &BoundVariableKinds) {
440                self.current_index.shift_in(1);
441            }
442
443            fn exit_binder(&mut self) {
444                self.current_index.shift_out(1);
445            }
446
447            fn fold_region(&mut self, re: &Region) -> Region {
448                if let ReBound(debruijn, br) = *re
449                    && debruijn == self.current_index
450                {
451                    ReBound(debruijn, BoundRegion { var: br.var + self.amount, kind: br.kind })
452                } else {
453                    *re
454                }
455            }
456
457            fn fold_expr(&mut self, expr: &Expr) -> Expr {
458                if let ExprKind::Var(Var::Bound(debruijn, breft)) = expr.kind()
459                    && debruijn == &self.current_index
460                {
461                    Expr::bvar(*debruijn, breft.var + self.amount, breft.kind)
462                } else {
463                    expr.super_fold_with(self)
464                }
465            }
466        }
467        self.fold_with(&mut Shifter { amount, current_index: INNERMOST })
468    }
469
470    fn shift_in_escaping(&self, amount: u32) -> Self {
471        struct Shifter {
472            current_index: DebruijnIndex,
473            amount: u32,
474        }
475
476        impl TypeFolder for Shifter {
477            fn enter_binder(&mut self, _: &BoundVariableKinds) {
478                self.current_index.shift_in(1);
479            }
480
481            fn exit_binder(&mut self) {
482                self.current_index.shift_out(1);
483            }
484
485            fn fold_region(&mut self, re: &Region) -> Region {
486                if let ReBound(debruijn, br) = *re
487                    && debruijn >= self.current_index
488                {
489                    ReBound(debruijn.shifted_in(self.amount), br)
490                } else {
491                    *re
492                }
493            }
494
495            fn fold_expr(&mut self, expr: &Expr) -> Expr {
496                if let ExprKind::Var(Var::Bound(debruijn, breft)) = expr.kind()
497                    && *debruijn >= self.current_index
498                {
499                    Expr::bvar(debruijn.shifted_in(self.amount), breft.var, breft.kind)
500                } else {
501                    expr.super_fold_with(self)
502                }
503            }
504        }
505        self.fold_with(&mut Shifter { amount, current_index: INNERMOST })
506    }
507
508    fn shift_out_escaping(&self, amount: u32) -> Self {
509        struct Shifter {
510            amount: u32,
511            current_index: DebruijnIndex,
512        }
513
514        impl TypeFolder for Shifter {
515            fn enter_binder(&mut self, _: &BoundVariableKinds) {
516                self.current_index.shift_in(1);
517            }
518
519            fn exit_binder(&mut self) {
520                self.current_index.shift_out(1);
521            }
522
523            fn fold_region(&mut self, re: &Region) -> Region {
524                if let ReBound(debruijn, br) = *re
525                    && debruijn >= self.current_index
526                {
527                    ReBound(debruijn.shifted_out(self.amount), br)
528                } else {
529                    *re
530                }
531            }
532
533            fn fold_expr(&mut self, expr: &Expr) -> Expr {
534                if let ExprKind::Var(Var::Bound(debruijn, breft)) = expr.kind()
535                    && debruijn >= &self.current_index
536                {
537                    Expr::bvar(debruijn.shifted_out(self.amount), breft.var, breft.kind)
538                } else {
539                    expr.super_fold_with(self)
540                }
541            }
542        }
543        self.fold_with(&mut Shifter { amount, current_index: INNERMOST })
544    }
545
546    fn erase_regions(&self) -> Self {
547        struct RegionEraser;
548        impl TypeFolder for RegionEraser {
549            fn fold_region(&mut self, r: &Region) -> Region {
550                match *r {
551                    ReBound(..) => *r,
552                    _ => ReErased,
553                }
554            }
555        }
556
557        self.fold_with(&mut RegionEraser)
558    }
559}
560
561pub trait TypeSuperFoldable: TypeFoldable {
562    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error>;
563
564    fn super_fold_with<F: TypeFolder>(&self, folder: &mut F) -> Self {
565        self.try_super_fold_with(folder).into_ok()
566    }
567}
568
569impl<T: TypeVisitable> TypeVisitable for OutlivesPredicate<T> {
570    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
571        self.0.visit_with(visitor)?;
572        self.1.visit_with(visitor)
573    }
574}
575
576impl<T: TypeFoldable> TypeFoldable for OutlivesPredicate<T> {
577    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
578        Ok(OutlivesPredicate(self.0.try_fold_with(folder)?, self.1.try_fold_with(folder)?))
579    }
580}
581
582impl TypeVisitable for Sort {
583    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
584        visitor.visit_sort(self)
585    }
586}
587
588impl TypeSuperVisitable for Sort {
589    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
590        match self {
591            Sort::Tuple(sorts) => sorts.visit_with(visitor),
592            Sort::App(_, args) => args.visit_with(visitor),
593            Sort::Func(fsort) => fsort.visit_with(visitor),
594            Sort::Alias(_, alias_ty) => alias_ty.visit_with(visitor),
595            Sort::Int
596            | Sort::Bool
597            | Sort::Real
598            | Sort::Str
599            | Sort::Char
600            | Sort::BitVec(_)
601            | Sort::Loc
602            | Sort::Param(_)
603            | Sort::Var(_)
604            | Sort::Infer(_)
605            | Sort::Err => ControlFlow::Continue(()),
606        }
607    }
608}
609
610impl TypeFoldable for Sort {
611    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
612        folder.try_fold_sort(self)
613    }
614}
615
616impl TypeSuperFoldable for Sort {
617    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
618        let sort = match self {
619            Sort::Tuple(sorts) => Sort::tuple(sorts.try_fold_with(folder)?),
620            Sort::App(ctor, sorts) => Sort::app(ctor.clone(), sorts.try_fold_with(folder)?),
621            Sort::Func(fsort) => Sort::Func(fsort.try_fold_with(folder)?),
622            Sort::Alias(kind, alias_ty) => Sort::Alias(*kind, alias_ty.try_fold_with(folder)?),
623            Sort::Int
624            | Sort::Bool
625            | Sort::Real
626            | Sort::Loc
627            | Sort::Str
628            | Sort::Char
629            | Sort::BitVec(_)
630            | Sort::Param(_)
631            | Sort::Var(_)
632            | Sort::Infer(_)
633            | Sort::Err => self.clone(),
634        };
635        Ok(sort)
636    }
637}
638
639impl TypeVisitable for PolyFuncSort {
640    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
641        self.fsort.visit_with(visitor)
642    }
643}
644
645impl TypeFoldable for PolyFuncSort {
646    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
647        Ok(PolyFuncSort { params: self.params.clone(), fsort: self.fsort.try_fold_with(folder)? })
648    }
649}
650
651impl<T> TypeVisitable for Binder<T>
652where
653    T: TypeVisitable,
654{
655    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
656        self.vars().visit_with(visitor)?;
657        visitor.enter_binder(self.vars());
658        self.skip_binder_ref().visit_with(visitor)?;
659        visitor.exit_binder();
660        ControlFlow::Continue(())
661    }
662}
663
664impl<T> TypeSuperVisitable for Binder<T>
665where
666    T: TypeVisitable,
667{
668    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
669        self.skip_binder_ref().visit_with(visitor)
670    }
671}
672
673impl<T> TypeFoldable for Binder<T>
674where
675    T: TypeFoldable,
676{
677    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
678        let vars = self.vars().try_fold_with(folder)?;
679        folder.try_enter_binder(&vars);
680        let r = self.skip_binder_ref().try_fold_with(folder)?;
681        folder.try_exit_binder();
682        Ok(Binder::bind_with_vars(r, vars))
683    }
684}
685
686impl TypeVisitable for VariantSig {
687    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
688        self.fields.visit_with(visitor)?;
689        self.idx.visit_with(visitor)
690    }
691}
692
693impl TypeFoldable for VariantSig {
694    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
695        let args = self.args.try_fold_with(folder)?;
696        let fields = self.fields.try_fold_with(folder)?;
697        let idx = self.idx.try_fold_with(folder)?;
698        let requires = self.requires.try_fold_with(folder)?;
699        Ok(VariantSig::new(self.adt_def.clone(), args, fields, idx, requires))
700    }
701}
702
703impl<T: TypeVisitable> TypeVisitable for Vec<T> {
704    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
705        self.iter().try_for_each(|t| t.visit_with(visitor))
706    }
707}
708
709impl TypeVisitable for Ensures {
710    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
711        match self {
712            Ensures::Type(path, ty) => {
713                path.to_expr().visit_with(visitor)?;
714                ty.visit_with(visitor)
715            }
716            Ensures::Pred(e) => e.visit_with(visitor),
717        }
718    }
719}
720
721impl TypeFoldable for Ensures {
722    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
723        let c = match self {
724            Ensures::Type(path, ty) => {
725                let path_expr = path.to_expr().try_fold_with(folder)?;
726                Ensures::Type(
727                    path_expr.to_path().unwrap_or_else(|| {
728                        bug!("invalid path `{path_expr:?}` produced when folding `{self:?}`",)
729                    }),
730                    ty.try_fold_with(folder)?,
731                )
732            }
733            Ensures::Pred(e) => Ensures::Pred(e.try_fold_with(folder)?),
734        };
735        Ok(c)
736    }
737}
738
739impl TypeVisitable for super::TyOrBase {
740    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
741        match self {
742            Self::Ty(ty) => ty.visit_with(visitor),
743            Self::Base(bty) => bty.visit_with(visitor),
744        }
745    }
746}
747
748impl TypeVisitable for Ty {
749    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
750        visitor.visit_ty(self)
751    }
752}
753
754impl TypeSuperVisitable for Ty {
755    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
756        match self.kind() {
757            TyKind::Indexed(bty, idxs) => {
758                bty.visit_with(visitor)?;
759                idxs.visit_with(visitor)
760            }
761            TyKind::Exists(exists) => exists.visit_with(visitor),
762            TyKind::StrgRef(_, path, ty) => {
763                path.to_expr().visit_with(visitor)?;
764                ty.visit_with(visitor)
765            }
766            TyKind::Ptr(_, path) => path.to_expr().visit_with(visitor),
767            TyKind::Constr(pred, ty) => {
768                pred.visit_with(visitor)?;
769                ty.visit_with(visitor)
770            }
771            TyKind::Downcast(.., args, _, fields) => {
772                args.visit_with(visitor)?;
773                fields.visit_with(visitor)
774            }
775            TyKind::Blocked(ty) => ty.visit_with(visitor),
776            TyKind::Infer(_) | TyKind::Param(_) | TyKind::Discr(..) | TyKind::Uninit => {
777                ControlFlow::Continue(())
778            }
779        }
780    }
781}
782
783impl TypeFoldable for Ty {
784    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
785        folder.try_fold_ty(self)
786    }
787}
788
789impl TypeSuperFoldable for Ty {
790    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Ty, F::Error> {
791        let ty = match self.kind() {
792            TyKind::Indexed(bty, idxs) => {
793                Ty::indexed(bty.try_fold_with(folder)?, idxs.try_fold_with(folder)?)
794            }
795            TyKind::Exists(exists) => TyKind::Exists(exists.try_fold_with(folder)?).intern(),
796            TyKind::StrgRef(re, path, ty) => {
797                Ty::strg_ref(
798                    re.try_fold_with(folder)?,
799                    path.to_expr()
800                        .try_fold_with(folder)?
801                        .to_path()
802                        .expect("type folding produced an invalid path"),
803                    ty.try_fold_with(folder)?,
804                )
805            }
806            TyKind::Ptr(pk, path) => {
807                let pk = match pk {
808                    PtrKind::Mut(re) => PtrKind::Mut(re.try_fold_with(folder)?),
809                    PtrKind::Box => PtrKind::Box,
810                };
811                Ty::ptr(
812                    pk,
813                    path.to_expr()
814                        .try_fold_with(folder)?
815                        .to_path()
816                        .expect("type folding produced an invalid path"),
817                )
818            }
819            TyKind::Constr(pred, ty) => {
820                Ty::constr(pred.try_fold_with(folder)?, ty.try_fold_with(folder)?)
821            }
822            TyKind::Downcast(adt, args, ty, variant, fields) => {
823                Ty::downcast(
824                    adt.clone(),
825                    args.clone(),
826                    ty.clone(),
827                    *variant,
828                    fields.try_fold_with(folder)?,
829                )
830            }
831            TyKind::Blocked(ty) => Ty::blocked(ty.try_fold_with(folder)?),
832            TyKind::Infer(_) | TyKind::Param(_) | TyKind::Uninit | TyKind::Discr(..) => {
833                self.clone()
834            }
835        };
836        Ok(ty)
837    }
838}
839
840impl TypeVisitable for Region {
841    fn visit_with<V: TypeVisitor>(&self, _visitor: &mut V) -> ControlFlow<V::BreakTy> {
842        ControlFlow::Continue(())
843    }
844}
845
846impl TypeFoldable for Region {
847    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
848        folder.try_fold_region(self)
849    }
850}
851
852impl TypeSuperFoldable for Const {
853    fn try_super_fold_with<F: FallibleTypeFolder>(
854        &self,
855        _folder: &mut F,
856    ) -> Result<Self, F::Error> {
857        // FIXME(nilehmann) we are not folding the type in `ConstKind::Value` because it's a rustc::ty::Ty
858        Ok(self.clone())
859    }
860}
861
862impl TypeVisitable for Const {
863    fn visit_with<V: TypeVisitor>(&self, _visitor: &mut V) -> ControlFlow<V::BreakTy> {
864        ControlFlow::Continue(())
865    }
866}
867
868impl TypeFoldable for Const {
869    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
870        folder.try_fold_const(self)
871    }
872}
873
874impl TypeVisitable for BaseTy {
875    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
876        visitor.visit_bty(self)
877    }
878}
879
880impl TypeSuperVisitable for BaseTy {
881    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
882        match self {
883            BaseTy::Adt(_, args) => args.visit_with(visitor),
884            BaseTy::FnDef(_, args) => args.visit_with(visitor),
885            BaseTy::Slice(ty) => ty.visit_with(visitor),
886            BaseTy::RawPtr(ty, _) => ty.visit_with(visitor),
887            BaseTy::RawPtrMetadata(ty) => ty.visit_with(visitor),
888            BaseTy::Ref(_, ty, _) => ty.visit_with(visitor),
889            BaseTy::FnPtr(poly_fn_sig) => poly_fn_sig.visit_with(visitor),
890            BaseTy::Tuple(tys) => tys.visit_with(visitor),
891            BaseTy::Alias(_, alias_ty) => alias_ty.visit_with(visitor),
892            BaseTy::Array(ty, _) => ty.visit_with(visitor),
893            BaseTy::Coroutine(_, resume_ty, upvars) => {
894                resume_ty.visit_with(visitor)?;
895                upvars.visit_with(visitor)
896            }
897            BaseTy::Dynamic(exi_preds, _) => exi_preds.visit_with(visitor),
898            BaseTy::Int(_)
899            | BaseTy::Uint(_)
900            | BaseTy::Bool
901            | BaseTy::Float(_)
902            | BaseTy::Str
903            | BaseTy::Char
904            | BaseTy::Closure(..)
905            | BaseTy::Never
906            | BaseTy::Infer(_)
907            | BaseTy::Foreign(_)
908            | BaseTy::Param(_) => ControlFlow::Continue(()),
909        }
910    }
911}
912
913impl TypeFoldable for BaseTy {
914    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
915        folder.try_fold_bty(self)
916    }
917}
918
919impl TypeSuperFoldable for BaseTy {
920    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
921        let bty = match self {
922            BaseTy::Adt(adt_def, args) => BaseTy::adt(adt_def.clone(), args.try_fold_with(folder)?),
923            BaseTy::FnDef(def_id, args) => BaseTy::fn_def(*def_id, args.try_fold_with(folder)?),
924            BaseTy::Foreign(def_id) => BaseTy::Foreign(*def_id),
925            BaseTy::Slice(ty) => BaseTy::Slice(ty.try_fold_with(folder)?),
926            BaseTy::RawPtr(ty, mu) => BaseTy::RawPtr(ty.try_fold_with(folder)?, *mu),
927            BaseTy::RawPtrMetadata(ty) => BaseTy::RawPtrMetadata(ty.try_fold_with(folder)?),
928            BaseTy::Ref(re, ty, mutbl) => {
929                BaseTy::Ref(re.try_fold_with(folder)?, ty.try_fold_with(folder)?, *mutbl)
930            }
931            BaseTy::FnPtr(decl) => BaseTy::FnPtr(decl.try_fold_with(folder)?),
932            BaseTy::Tuple(tys) => BaseTy::Tuple(tys.try_fold_with(folder)?),
933            BaseTy::Alias(kind, alias_ty) => BaseTy::Alias(*kind, alias_ty.try_fold_with(folder)?),
934            BaseTy::Array(ty, c) => {
935                BaseTy::Array(ty.try_fold_with(folder)?, c.try_fold_with(folder)?)
936            }
937            BaseTy::Closure(did, args, gen_args) => {
938                BaseTy::Closure(*did, args.try_fold_with(folder)?, gen_args.clone())
939            }
940            BaseTy::Coroutine(did, resume_ty, args) => {
941                BaseTy::Coroutine(
942                    *did,
943                    resume_ty.try_fold_with(folder)?,
944                    args.try_fold_with(folder)?,
945                )
946            }
947            BaseTy::Dynamic(preds, region) => {
948                BaseTy::Dynamic(preds.try_fold_with(folder)?, region.try_fold_with(folder)?)
949            }
950            BaseTy::Int(_)
951            | BaseTy::Param(_)
952            | BaseTy::Uint(_)
953            | BaseTy::Bool
954            | BaseTy::Float(_)
955            | BaseTy::Str
956            | BaseTy::Char
957            | BaseTy::Infer(_)
958            | BaseTy::Never => self.clone(),
959        };
960        Ok(bty)
961    }
962}
963
964impl TypeVisitable for SubsetTy {
965    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
966        self.bty.visit_with(visitor)?;
967        self.idx.visit_with(visitor)?;
968        self.pred.visit_with(visitor)
969    }
970}
971impl TypeFoldable for TyOrBase {
972    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
973        match self {
974            Self::Ty(ty) => Ok(Self::Ty(ty.try_fold_with(folder)?)),
975            Self::Base(bty) => Ok(Self::Base(bty.try_fold_with(folder)?)),
976        }
977    }
978}
979
980impl TypeFoldable for SubsetTy {
981    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
982        folder.try_fold_subset_ty(self)
983    }
984}
985
986impl TypeSuperFoldable for SubsetTy {
987    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
988        Ok(SubsetTy {
989            bty: self.bty.try_fold_with(folder)?,
990            idx: self.idx.try_fold_with(folder)?,
991            pred: self.pred.try_fold_with(folder)?,
992        })
993    }
994}
995
996impl TypeVisitable for GenericArg {
997    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
998        match self {
999            GenericArg::Ty(ty) => ty.visit_with(visitor),
1000            GenericArg::Base(ty) => ty.visit_with(visitor),
1001            GenericArg::Lifetime(_) => ControlFlow::Continue(()),
1002            GenericArg::Const(_) => ControlFlow::Continue(()),
1003        }
1004    }
1005}
1006
1007impl TypeFoldable for GenericArg {
1008    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1009        let arg = match self {
1010            GenericArg::Ty(ty) => GenericArg::Ty(ty.try_fold_with(folder)?),
1011            GenericArg::Base(ctor) => GenericArg::Base(ctor.try_fold_with(folder)?),
1012            GenericArg::Lifetime(re) => GenericArg::Lifetime(re.try_fold_with(folder)?),
1013            GenericArg::Const(c) => GenericArg::Const(c.try_fold_with(folder)?),
1014        };
1015        Ok(arg)
1016    }
1017}
1018
1019impl TypeVisitable for Expr {
1020    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1021        visitor.visit_expr(self)
1022    }
1023}
1024
1025impl TypeSuperVisitable for Expr {
1026    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1027        match self.kind() {
1028            ExprKind::Var(_) => ControlFlow::Continue(()),
1029            ExprKind::BinaryOp(_, e1, e2) => {
1030                e1.visit_with(visitor)?;
1031                e2.visit_with(visitor)
1032            }
1033            ExprKind::Tuple(flds) => flds.visit_with(visitor),
1034            ExprKind::Ctor(_, flds) => flds.visit_with(visitor),
1035            ExprKind::IsCtor(_, _, e) => e.visit_with(visitor),
1036            ExprKind::FieldProj(e, _) | ExprKind::PathProj(e, _) | ExprKind::UnaryOp(_, e) => {
1037                e.visit_with(visitor)
1038            }
1039            ExprKind::App(func, sorts, arg) => {
1040                func.visit_with(visitor)?;
1041                sorts.visit_with(visitor)?;
1042                arg.visit_with(visitor)
1043            }
1044            ExprKind::IfThenElse(p, e1, e2) => {
1045                p.visit_with(visitor)?;
1046                e1.visit_with(visitor)?;
1047                e2.visit_with(visitor)
1048            }
1049            ExprKind::KVar(kvar) => kvar.visit_with(visitor),
1050            ExprKind::Alias(alias, args) => {
1051                alias.visit_with(visitor)?;
1052                args.visit_with(visitor)
1053            }
1054            ExprKind::Abs(body) => body.visit_with(visitor),
1055            ExprKind::BoundedQuant(_, _, body) => body.visit_with(visitor),
1056            ExprKind::ForAll(expr) => expr.visit_with(visitor),
1057            ExprKind::Exists(expr) => expr.visit_with(visitor),
1058            ExprKind::Let(init, body) => {
1059                init.visit_with(visitor)?;
1060                body.visit_with(visitor)
1061            }
1062            ExprKind::Constant(_)
1063            | ExprKind::Hole(_)
1064            | ExprKind::Local(_)
1065            | ExprKind::GlobalFunc(..)
1066            | ExprKind::InternalFunc(..)
1067            | ExprKind::ConstDefId(_) => ControlFlow::Continue(()),
1068        }
1069    }
1070}
1071
1072impl TypeFoldable for Expr {
1073    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1074        folder.try_fold_expr(self)
1075    }
1076}
1077
1078impl TypeSuperFoldable for Expr {
1079    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1080        let span = self.span();
1081        let expr = match self.kind() {
1082            ExprKind::Var(var) => Expr::var(*var),
1083            ExprKind::Local(local) => Expr::local(*local),
1084            ExprKind::Constant(c) => Expr::constant(*c),
1085            ExprKind::ConstDefId(did) => Expr::const_def_id(*did),
1086            ExprKind::BinaryOp(op, e1, e2) => {
1087                Expr::binary_op(
1088                    op.try_fold_with(folder)?,
1089                    e1.try_fold_with(folder)?,
1090                    e2.try_fold_with(folder)?,
1091                )
1092            }
1093            ExprKind::UnaryOp(op, e) => Expr::unary_op(*op, e.try_fold_with(folder)?),
1094            ExprKind::FieldProj(e, proj) => Expr::field_proj(e.try_fold_with(folder)?, *proj),
1095            ExprKind::Tuple(flds) => Expr::tuple(flds.try_fold_with(folder)?),
1096            ExprKind::Ctor(ctor, flds) => Expr::ctor(*ctor, flds.try_fold_with(folder)?),
1097            ExprKind::IsCtor(def_id, variant_idx, e) => {
1098                Expr::is_ctor(*def_id, *variant_idx, e.try_fold_with(folder)?)
1099            }
1100            ExprKind::PathProj(e, field) => Expr::path_proj(e.try_fold_with(folder)?, *field),
1101            ExprKind::App(func, sorts, arg) => {
1102                Expr::app(
1103                    func.try_fold_with(folder)?,
1104                    sorts.try_fold_with(folder)?,
1105                    arg.try_fold_with(folder)?,
1106                )
1107            }
1108            ExprKind::IfThenElse(p, e1, e2) => {
1109                Expr::ite(
1110                    p.try_fold_with(folder)?,
1111                    e1.try_fold_with(folder)?,
1112                    e2.try_fold_with(folder)?,
1113                )
1114            }
1115            ExprKind::Hole(kind) => Expr::hole(kind.try_fold_with(folder)?),
1116            ExprKind::KVar(kvar) => Expr::kvar(kvar.try_fold_with(folder)?),
1117            ExprKind::Abs(lam) => Expr::abs(lam.try_fold_with(folder)?),
1118            ExprKind::BoundedQuant(kind, rng, body) => {
1119                Expr::bounded_quant(*kind, *rng, body.try_fold_with(folder)?)
1120            }
1121            ExprKind::GlobalFunc(kind) => Expr::global_func(kind.clone()),
1122            ExprKind::InternalFunc(kind) => Expr::internal_func(kind.clone()),
1123            ExprKind::Alias(alias, args) => {
1124                Expr::alias(alias.try_fold_with(folder)?, args.try_fold_with(folder)?)
1125            }
1126            ExprKind::ForAll(expr) => Expr::forall(expr.try_fold_with(folder)?),
1127            ExprKind::Exists(expr) => Expr::exists(expr.try_fold_with(folder)?),
1128            ExprKind::Let(init, body) => {
1129                Expr::let_(init.try_fold_with(folder)?, body.try_fold_with(folder)?)
1130            }
1131        };
1132        Ok(expr.at_opt(span))
1133    }
1134}
1135
1136impl<T> TypeVisitable for List<T>
1137where
1138    T: TypeVisitable,
1139    [T]: Internable,
1140{
1141    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1142        self.iter().try_for_each(|t| t.visit_with(visitor))
1143    }
1144}
1145
1146impl<T> TypeFoldable for List<T>
1147where
1148    T: TypeFoldable,
1149    [T]: Internable,
1150{
1151    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1152        self.iter().map(|t| t.try_fold_with(folder)).try_collect()
1153    }
1154}
1155
1156/// Used for types that are `Copy` and which **do not care arena allocated data** (i.e., don't need
1157/// to be folded).
1158macro_rules! TrivialTypeTraversalImpls {
1159    ($($ty:ty,)+) => {
1160        $(
1161            impl $crate::rty::fold::TypeFoldable for $ty {
1162                fn try_fold_with<F: $crate::rty::fold::FallibleTypeFolder>(
1163                    &self,
1164                    _: &mut F,
1165                ) -> ::std::result::Result<Self, F::Error> {
1166                    Ok(*self)
1167                }
1168
1169                #[inline]
1170                fn fold_with<F: $crate::rty::fold::TypeFolder>(
1171                    &self,
1172                    _: &mut F,
1173                ) -> Self {
1174                    *self
1175                }
1176            }
1177
1178            impl $crate::rty::fold::TypeVisitable for $ty {
1179                #[inline]
1180                fn visit_with<V: $crate::rty::fold::TypeVisitor>(
1181                    &self,
1182                    _: &mut V)
1183                    -> ::core::ops::ControlFlow<V::BreakTy>
1184                {
1185                    ::core::ops::ControlFlow::Continue(())
1186                }
1187            }
1188        )+
1189    };
1190}
1191
1192// For things that don't carry any arena-allocated data (and are copy...), just add them to this list.
1193TrivialTypeTraversalImpls! {
1194    (),
1195    bool,
1196    usize,
1197    crate::fhir::InferMode,
1198    crate::rty::BoundReftKind,
1199    crate::rty::BvSize,
1200    crate::rty::KVid,
1201    crate::def_id::FluxDefId,
1202    crate::def_id::FluxLocalDefId,
1203    rustc_span::Symbol,
1204    rustc_hir::def_id::DefId,
1205    rustc_hir::Safety,
1206    rustc_abi::ExternAbi,
1207    rustc_type_ir::ClosureKind,
1208    flux_rustc_bridge::ty::BoundRegionKind,
1209}
flux_middle/rty/fold.rs

flux_middle/rty/
fold.rs
