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::RawPtr
606            | Sort::Err => ControlFlow::Continue(()),
607        }
608    }
609}
610
611impl TypeFoldable for Sort {
612    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
613        folder.try_fold_sort(self)
614    }
615}
616
617impl TypeSuperFoldable for Sort {
618    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
619        let sort = match self {
620            Sort::Tuple(sorts) => Sort::tuple(sorts.try_fold_with(folder)?),
621            Sort::App(ctor, sorts) => Sort::app(ctor.clone(), sorts.try_fold_with(folder)?),
622            Sort::Func(fsort) => Sort::Func(fsort.try_fold_with(folder)?),
623            Sort::Alias(kind, alias_ty) => Sort::Alias(*kind, alias_ty.try_fold_with(folder)?),
624            Sort::Int
625            | Sort::Bool
626            | Sort::Real
627            | Sort::Loc
628            | Sort::Str
629            | Sort::Char
630            | Sort::BitVec(_)
631            | Sort::Param(_)
632            | Sort::Var(_)
633            | Sort::Infer(_)
634            | Sort::RawPtr
635            | Sort::Err => self.clone(),
636        };
637        Ok(sort)
638    }
639}
640
641impl TypeVisitable for PolyFuncSort {
642    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
643        self.fsort.visit_with(visitor)
644    }
645}
646
647impl TypeFoldable for PolyFuncSort {
648    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
649        Ok(PolyFuncSort { params: self.params.clone(), fsort: self.fsort.try_fold_with(folder)? })
650    }
651}
652
653impl<T> TypeVisitable for Binder<T>
654where
655    T: TypeVisitable,
656{
657    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
658        self.vars().visit_with(visitor)?;
659        visitor.enter_binder(self.vars());
660        self.skip_binder_ref().visit_with(visitor)?;
661        visitor.exit_binder();
662        ControlFlow::Continue(())
663    }
664}
665
666impl<T> TypeSuperVisitable for Binder<T>
667where
668    T: TypeVisitable,
669{
670    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
671        self.skip_binder_ref().visit_with(visitor)
672    }
673}
674
675impl<T> TypeFoldable for Binder<T>
676where
677    T: TypeFoldable,
678{
679    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
680        let vars = self.vars().try_fold_with(folder)?;
681        folder.try_enter_binder(&vars);
682        let r = self.skip_binder_ref().try_fold_with(folder)?;
683        folder.try_exit_binder();
684        Ok(Binder::bind_with_vars(r, vars))
685    }
686}
687
688impl TypeVisitable for VariantSig {
689    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
690        self.fields.visit_with(visitor)?;
691        self.idx.visit_with(visitor)
692    }
693}
694
695impl TypeFoldable for VariantSig {
696    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
697        let args = self.args.try_fold_with(folder)?;
698        let fields = self.fields.try_fold_with(folder)?;
699        let idx = self.idx.try_fold_with(folder)?;
700        let requires = self.requires.try_fold_with(folder)?;
701        Ok(VariantSig::new(self.adt_def.clone(), args, fields, idx, requires))
702    }
703}
704
705impl<T: TypeVisitable> TypeVisitable for Vec<T> {
706    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
707        self.iter().try_for_each(|t| t.visit_with(visitor))
708    }
709}
710
711impl TypeVisitable for Ensures {
712    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
713        match self {
714            Ensures::Type(path, ty) => {
715                path.to_expr().visit_with(visitor)?;
716                ty.visit_with(visitor)
717            }
718            Ensures::Pred(e) => e.visit_with(visitor),
719        }
720    }
721}
722
723impl TypeFoldable for Ensures {
724    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
725        let c = match self {
726            Ensures::Type(path, ty) => {
727                let path_expr = path.to_expr().try_fold_with(folder)?;
728                Ensures::Type(
729                    path_expr.to_path().unwrap_or_else(|| {
730                        bug!("invalid path `{path_expr:?}` produced when folding `{self:?}`",)
731                    }),
732                    ty.try_fold_with(folder)?,
733                )
734            }
735            Ensures::Pred(e) => Ensures::Pred(e.try_fold_with(folder)?),
736        };
737        Ok(c)
738    }
739}
740
741impl TypeVisitable for super::TyOrBase {
742    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
743        match self {
744            Self::Ty(ty) => ty.visit_with(visitor),
745            Self::Base(bty) => bty.visit_with(visitor),
746        }
747    }
748}
749
750impl TypeVisitable for Ty {
751    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
752        visitor.visit_ty(self)
753    }
754}
755
756impl TypeSuperVisitable for Ty {
757    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
758        match self.kind() {
759            TyKind::Indexed(bty, idxs) => {
760                bty.visit_with(visitor)?;
761                idxs.visit_with(visitor)
762            }
763            TyKind::Exists(exists) => exists.visit_with(visitor),
764            TyKind::StrgRef(_, path, ty) => {
765                path.to_expr().visit_with(visitor)?;
766                ty.visit_with(visitor)
767            }
768            TyKind::Ptr(_, path) => path.to_expr().visit_with(visitor),
769            TyKind::Constr(pred, ty) => {
770                pred.visit_with(visitor)?;
771                ty.visit_with(visitor)
772            }
773            TyKind::Downcast(.., args, _, fields) => {
774                args.visit_with(visitor)?;
775                fields.visit_with(visitor)
776            }
777            TyKind::Blocked(ty) => ty.visit_with(visitor),
778            TyKind::Infer(_) | TyKind::Param(_) | TyKind::Discr(..) | TyKind::Uninit => {
779                ControlFlow::Continue(())
780            }
781        }
782    }
783}
784
785impl TypeFoldable for Ty {
786    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
787        folder.try_fold_ty(self)
788    }
789}
790
791impl TypeSuperFoldable for Ty {
792    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Ty, F::Error> {
793        let ty = match self.kind() {
794            TyKind::Indexed(bty, idxs) => {
795                Ty::indexed(bty.try_fold_with(folder)?, idxs.try_fold_with(folder)?)
796            }
797            TyKind::Exists(exists) => TyKind::Exists(exists.try_fold_with(folder)?).intern(),
798            TyKind::StrgRef(re, path, ty) => {
799                Ty::strg_ref(
800                    re.try_fold_with(folder)?,
801                    path.to_expr()
802                        .try_fold_with(folder)?
803                        .to_path()
804                        .expect("type folding produced an invalid path"),
805                    ty.try_fold_with(folder)?,
806                )
807            }
808            TyKind::Ptr(pk, path) => {
809                let pk = match pk {
810                    PtrKind::Mut(re) => PtrKind::Mut(re.try_fold_with(folder)?),
811                    PtrKind::Box => PtrKind::Box,
812                };
813                Ty::ptr(
814                    pk,
815                    path.to_expr()
816                        .try_fold_with(folder)?
817                        .to_path()
818                        .expect("type folding produced an invalid path"),
819                )
820            }
821            TyKind::Constr(pred, ty) => {
822                Ty::constr(pred.try_fold_with(folder)?, ty.try_fold_with(folder)?)
823            }
824            TyKind::Downcast(adt, args, ty, variant, fields) => {
825                Ty::downcast(
826                    adt.clone(),
827                    args.clone(),
828                    ty.clone(),
829                    *variant,
830                    fields.try_fold_with(folder)?,
831                )
832            }
833            TyKind::Blocked(ty) => Ty::blocked(ty.try_fold_with(folder)?),
834            TyKind::Infer(_) | TyKind::Param(_) | TyKind::Uninit | TyKind::Discr(..) => {
835                self.clone()
836            }
837        };
838        Ok(ty)
839    }
840}
841
842impl TypeVisitable for Region {
843    fn visit_with<V: TypeVisitor>(&self, _visitor: &mut V) -> ControlFlow<V::BreakTy> {
844        ControlFlow::Continue(())
845    }
846}
847
848impl TypeFoldable for Region {
849    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
850        folder.try_fold_region(self)
851    }
852}
853
854impl TypeSuperFoldable for Const {
855    fn try_super_fold_with<F: FallibleTypeFolder>(
856        &self,
857        _folder: &mut F,
858    ) -> Result<Self, F::Error> {
859        // FIXME(nilehmann) we are not folding the type in `ConstKind::Value` because it's a rustc::ty::Ty
860        Ok(self.clone())
861    }
862}
863
864impl TypeVisitable for Const {
865    fn visit_with<V: TypeVisitor>(&self, _visitor: &mut V) -> ControlFlow<V::BreakTy> {
866        ControlFlow::Continue(())
867    }
868}
869
870impl TypeFoldable for Const {
871    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
872        folder.try_fold_const(self)
873    }
874}
875
876impl TypeVisitable for BaseTy {
877    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
878        visitor.visit_bty(self)
879    }
880}
881
882impl TypeSuperVisitable for BaseTy {
883    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
884        match self {
885            BaseTy::Adt(_, args) => args.visit_with(visitor),
886            BaseTy::FnDef(_, args) => args.visit_with(visitor),
887            BaseTy::Slice(ty) => ty.visit_with(visitor),
888            BaseTy::RawPtr(ty, _) => ty.visit_with(visitor),
889            BaseTy::RawPtrMetadata(ty) => ty.visit_with(visitor),
890            BaseTy::Ref(_, ty, _) => ty.visit_with(visitor),
891            BaseTy::FnPtr(poly_fn_sig) => poly_fn_sig.visit_with(visitor),
892            BaseTy::Tuple(tys) => tys.visit_with(visitor),
893            BaseTy::Alias(_, alias_ty) => alias_ty.visit_with(visitor),
894            BaseTy::Array(ty, _) => ty.visit_with(visitor),
895            BaseTy::Coroutine(_, resume_ty, upvars, _) => {
896                resume_ty.visit_with(visitor)?;
897                upvars.visit_with(visitor)
898            }
899            BaseTy::Dynamic(exi_preds, _) => exi_preds.visit_with(visitor),
900            BaseTy::Int(_)
901            | BaseTy::Pat
902            | BaseTy::Uint(_)
903            | BaseTy::Bool
904            | BaseTy::Float(_)
905            | BaseTy::Str
906            | BaseTy::Char
907            | BaseTy::Closure(..)
908            | BaseTy::Never
909            | BaseTy::Infer(_)
910            | BaseTy::Foreign(_)
911            | BaseTy::Param(_) => ControlFlow::Continue(()),
912        }
913    }
914}
915
916impl TypeFoldable for BaseTy {
917    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
918        folder.try_fold_bty(self)
919    }
920}
921
922impl TypeSuperFoldable for BaseTy {
923    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
924        let bty = match self {
925            BaseTy::Adt(adt_def, args) => BaseTy::adt(adt_def.clone(), args.try_fold_with(folder)?),
926            BaseTy::FnDef(def_id, args) => BaseTy::fn_def(*def_id, args.try_fold_with(folder)?),
927            BaseTy::Foreign(def_id) => BaseTy::Foreign(*def_id),
928            BaseTy::Slice(ty) => BaseTy::Slice(ty.try_fold_with(folder)?),
929            BaseTy::RawPtr(ty, mu) => BaseTy::RawPtr(ty.try_fold_with(folder)?, *mu),
930            BaseTy::RawPtrMetadata(ty) => BaseTy::RawPtrMetadata(ty.try_fold_with(folder)?),
931            BaseTy::Ref(re, ty, mutbl) => {
932                BaseTy::Ref(re.try_fold_with(folder)?, ty.try_fold_with(folder)?, *mutbl)
933            }
934            BaseTy::FnPtr(decl) => BaseTy::FnPtr(decl.try_fold_with(folder)?),
935            BaseTy::Tuple(tys) => BaseTy::Tuple(tys.try_fold_with(folder)?),
936            BaseTy::Alias(kind, alias_ty) => BaseTy::Alias(*kind, alias_ty.try_fold_with(folder)?),
937            BaseTy::Array(ty, c) => {
938                BaseTy::Array(ty.try_fold_with(folder)?, c.try_fold_with(folder)?)
939            }
940            BaseTy::Closure(did, args, gen_args, no_panic) => {
941                BaseTy::Closure(*did, args.try_fold_with(folder)?, gen_args.clone(), *no_panic)
942            }
943            BaseTy::Coroutine(did, resume_ty, args, gen_args) => {
944                BaseTy::Coroutine(
945                    *did,
946                    resume_ty.try_fold_with(folder)?,
947                    args.try_fold_with(folder)?,
948                    gen_args.clone(),
949                )
950            }
951            BaseTy::Dynamic(preds, region) => {
952                BaseTy::Dynamic(preds.try_fold_with(folder)?, region.try_fold_with(folder)?)
953            }
954            BaseTy::Pat => BaseTy::Pat,
955            BaseTy::Int(_)
956            | BaseTy::Param(_)
957            | BaseTy::Uint(_)
958            | BaseTy::Bool
959            | BaseTy::Float(_)
960            | BaseTy::Str
961            | BaseTy::Char
962            | BaseTy::Infer(_)
963            | BaseTy::Never => self.clone(),
964        };
965        Ok(bty)
966    }
967}
968
969impl TypeVisitable for SubsetTy {
970    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
971        self.bty.visit_with(visitor)?;
972        self.idx.visit_with(visitor)?;
973        self.pred.visit_with(visitor)
974    }
975}
976impl TypeFoldable for TyOrBase {
977    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
978        match self {
979            Self::Ty(ty) => Ok(Self::Ty(ty.try_fold_with(folder)?)),
980            Self::Base(bty) => Ok(Self::Base(bty.try_fold_with(folder)?)),
981        }
982    }
983}
984
985impl TypeFoldable for SubsetTy {
986    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
987        folder.try_fold_subset_ty(self)
988    }
989}
990
991impl TypeSuperFoldable for SubsetTy {
992    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
993        Ok(SubsetTy {
994            bty: self.bty.try_fold_with(folder)?,
995            idx: self.idx.try_fold_with(folder)?,
996            pred: self.pred.try_fold_with(folder)?,
997        })
998    }
999}
1000
1001impl TypeVisitable for GenericArg {
1002    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1003        match self {
1004            GenericArg::Ty(ty) => ty.visit_with(visitor),
1005            GenericArg::Base(ty) => ty.visit_with(visitor),
1006            GenericArg::Lifetime(_) => ControlFlow::Continue(()),
1007            GenericArg::Const(_) => ControlFlow::Continue(()),
1008        }
1009    }
1010}
1011
1012impl TypeFoldable for GenericArg {
1013    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1014        let arg = match self {
1015            GenericArg::Ty(ty) => GenericArg::Ty(ty.try_fold_with(folder)?),
1016            GenericArg::Base(ctor) => GenericArg::Base(ctor.try_fold_with(folder)?),
1017            GenericArg::Lifetime(re) => GenericArg::Lifetime(re.try_fold_with(folder)?),
1018            GenericArg::Const(c) => GenericArg::Const(c.try_fold_with(folder)?),
1019        };
1020        Ok(arg)
1021    }
1022}
1023
1024impl TypeVisitable for Expr {
1025    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1026        visitor.visit_expr(self)
1027    }
1028}
1029
1030impl TypeSuperVisitable for Expr {
1031    fn super_visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1032        match self.kind() {
1033            ExprKind::Var(_) => ControlFlow::Continue(()),
1034            ExprKind::BinaryOp(_, e1, e2) => {
1035                e1.visit_with(visitor)?;
1036                e2.visit_with(visitor)
1037            }
1038            ExprKind::Tuple(flds) => flds.visit_with(visitor),
1039            ExprKind::Ctor(_, flds) => flds.visit_with(visitor),
1040            ExprKind::IsCtor(_, _, e) => e.visit_with(visitor),
1041            ExprKind::FieldProj(e, _) | ExprKind::PathProj(e, _) | ExprKind::UnaryOp(_, e) => {
1042                e.visit_with(visitor)
1043            }
1044            ExprKind::App(func, sorts, arg) => {
1045                func.visit_with(visitor)?;
1046                sorts.visit_with(visitor)?;
1047                arg.visit_with(visitor)
1048            }
1049            ExprKind::IfThenElse(p, e1, e2) => {
1050                p.visit_with(visitor)?;
1051                e1.visit_with(visitor)?;
1052                e2.visit_with(visitor)
1053            }
1054            ExprKind::KVar(kvar) => kvar.visit_with(visitor),
1055            ExprKind::Alias(alias, args) => {
1056                alias.visit_with(visitor)?;
1057                args.visit_with(visitor)
1058            }
1059            ExprKind::Abs(body) => body.visit_with(visitor),
1060            ExprKind::BoundedQuant(_, _, body) => body.visit_with(visitor),
1061            ExprKind::ForAll(expr) => expr.visit_with(visitor),
1062            ExprKind::Exists(expr) => expr.visit_with(visitor),
1063            ExprKind::Let(init, body) => {
1064                init.visit_with(visitor)?;
1065                body.visit_with(visitor)
1066            }
1067            ExprKind::Constant(_)
1068            | ExprKind::Hole(_)
1069            | ExprKind::Local(_)
1070            | ExprKind::GlobalFunc(..)
1071            | ExprKind::InternalFunc(..)
1072            | ExprKind::ConstDefId(_) => ControlFlow::Continue(()),
1073        }
1074    }
1075}
1076
1077impl TypeFoldable for Expr {
1078    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1079        folder.try_fold_expr(self)
1080    }
1081}
1082
1083impl TypeSuperFoldable for Expr {
1084    fn try_super_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1085        let span = self.span();
1086        let expr = match self.kind() {
1087            ExprKind::Var(var) => Expr::var(*var),
1088            ExprKind::Local(local) => Expr::local(*local),
1089            ExprKind::Constant(c) => Expr::constant(*c),
1090            ExprKind::ConstDefId(did) => Expr::const_def_id(*did),
1091            ExprKind::BinaryOp(op, e1, e2) => {
1092                Expr::binary_op(
1093                    op.try_fold_with(folder)?,
1094                    e1.try_fold_with(folder)?,
1095                    e2.try_fold_with(folder)?,
1096                )
1097            }
1098            ExprKind::UnaryOp(op, e) => Expr::unary_op(*op, e.try_fold_with(folder)?),
1099            ExprKind::FieldProj(e, proj) => Expr::field_proj(e.try_fold_with(folder)?, *proj),
1100            ExprKind::Tuple(flds) => Expr::tuple(flds.try_fold_with(folder)?),
1101            ExprKind::Ctor(ctor, flds) => Expr::ctor(*ctor, flds.try_fold_with(folder)?),
1102            ExprKind::IsCtor(def_id, variant_idx, e) => {
1103                Expr::is_ctor(*def_id, *variant_idx, e.try_fold_with(folder)?)
1104            }
1105            ExprKind::PathProj(e, field) => Expr::path_proj(e.try_fold_with(folder)?, *field),
1106            ExprKind::App(func, sorts, arg) => {
1107                Expr::app(
1108                    func.try_fold_with(folder)?,
1109                    sorts.try_fold_with(folder)?,
1110                    arg.try_fold_with(folder)?,
1111                )
1112            }
1113            ExprKind::IfThenElse(p, e1, e2) => {
1114                Expr::ite(
1115                    p.try_fold_with(folder)?,
1116                    e1.try_fold_with(folder)?,
1117                    e2.try_fold_with(folder)?,
1118                )
1119            }
1120            ExprKind::Hole(kind) => Expr::hole(kind.try_fold_with(folder)?),
1121            ExprKind::KVar(kvar) => Expr::kvar(kvar.try_fold_with(folder)?),
1122            ExprKind::Abs(lam) => Expr::abs(lam.try_fold_with(folder)?),
1123            ExprKind::BoundedQuant(kind, rng, body) => {
1124                Expr::bounded_quant(*kind, *rng, body.try_fold_with(folder)?)
1125            }
1126            ExprKind::GlobalFunc(kind) => Expr::global_func(kind.clone()),
1127            ExprKind::InternalFunc(kind) => Expr::internal_func(kind.clone()),
1128            ExprKind::Alias(alias, args) => {
1129                Expr::alias(alias.try_fold_with(folder)?, args.try_fold_with(folder)?)
1130            }
1131            ExprKind::ForAll(expr) => Expr::forall(expr.try_fold_with(folder)?),
1132            ExprKind::Exists(expr) => Expr::exists(expr.try_fold_with(folder)?),
1133            ExprKind::Let(init, body) => {
1134                Expr::let_(init.try_fold_with(folder)?, body.try_fold_with(folder)?)
1135            }
1136        };
1137        Ok(expr.at_opt(span))
1138    }
1139}
1140
1141impl<T> TypeVisitable for List<T>
1142where
1143    T: TypeVisitable,
1144    [T]: Internable,
1145{
1146    fn visit_with<V: TypeVisitor>(&self, visitor: &mut V) -> ControlFlow<V::BreakTy> {
1147        self.iter().try_for_each(|t| t.visit_with(visitor))
1148    }
1149}
1150
1151impl<T> TypeFoldable for List<T>
1152where
1153    T: TypeFoldable,
1154    [T]: Internable,
1155{
1156    fn try_fold_with<F: FallibleTypeFolder>(&self, folder: &mut F) -> Result<Self, F::Error> {
1157        self.iter().map(|t| t.try_fold_with(folder)).try_collect()
1158    }
1159}
1160
1161/// Used for types that are `Copy` and which **do not care arena allocated data** (i.e., don't need
1162/// to be folded).
1163macro_rules! TrivialTypeTraversalImpls {
1164    ($($ty:ty,)+) => {
1165        $(
1166            impl $crate::rty::fold::TypeFoldable for $ty {
1167                fn try_fold_with<F: $crate::rty::fold::FallibleTypeFolder>(
1168                    &self,
1169                    _: &mut F,
1170                ) -> ::std::result::Result<Self, F::Error> {
1171                    Ok(*self)
1172                }
1173
1174                #[inline]
1175                fn fold_with<F: $crate::rty::fold::TypeFolder>(
1176                    &self,
1177                    _: &mut F,
1178                ) -> Self {
1179                    *self
1180                }
1181            }
1182
1183            impl $crate::rty::fold::TypeVisitable for $ty {
1184                #[inline]
1185                fn visit_with<V: $crate::rty::fold::TypeVisitor>(
1186                    &self,
1187                    _: &mut V)
1188                    -> ::core::ops::ControlFlow<V::BreakTy>
1189                {
1190                    ::core::ops::ControlFlow::Continue(())
1191                }
1192            }
1193        )+
1194    };
1195}
1196
1197// For things that don't carry any arena-allocated data (and are copy...), just add them to this list.
1198TrivialTypeTraversalImpls! {
1199    (),
1200    bool,
1201    usize,
1202    crate::fhir::InferMode,
1203    crate::rty::BoundReftKind,
1204    crate::rty::BvSize,
1205    crate::rty::KVid,
1206    crate::def_id::FluxDefId,
1207    crate::def_id::FluxLocalDefId,
1208    rustc_span::Symbol,
1209    rustc_hir::def_id::DefId,
1210    rustc_hir::Safety,
1211    rustc_abi::ExternAbi,
1212    rustc_type_ir::ClosureKind,
1213    flux_rustc_bridge::ty::BoundRegionKind,
1214}
flux_middle/rty/fold.rs

flux_middle/rty/
fold.rs
