flux_middle/rty/
fold.rs

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