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

flux_middle/rty/
fold.rs
