rustc_type_ir/
fold.rs

1//! A folding traversal mechanism for complex data structures that contain type
2//! information.
3//!
4//! This is a modifying traversal. It consumes the data structure, producing a
5//! (possibly) modified version of it. Both fallible and infallible versions are
6//! available. The name is potentially confusing, because this traversal is more
7//! like `Iterator::map` than `Iterator::fold`.
8//!
9//! This traversal has limited flexibility. Only a small number of "types of
10//! interest" within the complex data structures can receive custom
11//! modification. These are the ones containing the most important type-related
12//! information, such as `Ty`, `Predicate`, `Region`, and `Const`.
13//!
14//! There are three traits involved in each traversal.
15//! - `TypeFoldable`. This is implemented once for many types, including:
16//!   - Types of interest, for which the methods delegate to the folder.
17//!   - All other types, including generic containers like `Vec` and `Option`.
18//!     It defines a "skeleton" of how they should be folded.
19//! - `TypeSuperFoldable`. This is implemented only for recursive types of
20//!   interest, and defines the folding "skeleton" for these types. (This
21//!   excludes `Region` because it is non-recursive, i.e. it never contains
22//!   other types of interest.)
23//! - `TypeFolder`/`FallibleTypeFolder`. One of these is implemented for each
24//!   folder. This defines how types of interest are folded.
25//!
26//! This means each fold is a mixture of (a) generic folding operations, and (b)
27//! custom fold operations that are specific to the folder.
28//! - The `TypeFoldable` impls handle most of the traversal, and call into
29//!   `TypeFolder`/`FallibleTypeFolder` when they encounter a type of interest.
30//! - A `TypeFolder`/`FallibleTypeFolder` may call into another `TypeFoldable`
31//!   impl, because some of the types of interest are recursive and can contain
32//!   other types of interest.
33//! - A `TypeFolder`/`FallibleTypeFolder` may also call into a `TypeSuperFoldable`
34//!   impl, because each folder might provide custom handling only for some types
35//!   of interest, or only for some variants of each type of interest, and then
36//!   use default traversal for the remaining cases.
37//!
38//! For example, if you have `struct S(Ty, U)` where `S: TypeFoldable` and `U:
39//! TypeFoldable`, and an instance `s = S(ty, u)`, it would be folded like so:
40//! ```text
41//! s.fold_with(folder) calls
42//! - ty.fold_with(folder) calls
43//!   - folder.fold_ty(ty) may call
44//!     - ty.super_fold_with(folder)
45//! - u.fold_with(folder)
46//! ```
47
48use std::convert::Infallible;
49use std::mem;
50use std::sync::Arc;
51
52use rustc_index::{Idx, IndexVec};
53use thin_vec::ThinVec;
54use tracing::{debug, instrument};
55
56use crate::inherent::*;
57use crate::visit::{TypeVisitable, TypeVisitableExt as _};
58use crate::{self as ty, Interner, TypeFlags};
59
60/// This trait is implemented for every type that can be folded,
61/// providing the skeleton of the traversal.
62///
63/// To implement this conveniently, use the derive macro located in
64/// `rustc_macros`.
65///
66/// This trait is a sub-trait of `TypeVisitable`. This is because many
67/// `TypeFolder` instances use the methods in `TypeVisitableExt` while folding,
68/// which means in practice almost every foldable type needs to also be
69/// visitable. (However, there are some types that are visitable without being
70/// foldable.)
71pub trait TypeFoldable<I: Interner>: TypeVisitable<I> + Clone {
72    /// The entry point for folding. To fold a value `t` with a folder `f`
73    /// call: `t.try_fold_with(f)`.
74    ///
75    /// For most types, this just traverses the value, calling `try_fold_with`
76    /// on each field/element.
77    ///
78    /// For types of interest (such as `Ty`), the implementation of this method
79    /// calls a folder method specifically for that type (such as
80    /// `F::try_fold_ty`). This is where control transfers from [`TypeFoldable`]
81    /// to [`FallibleTypeFolder`].
82    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error>;
83
84    /// The entry point for folding. To fold a value `t` with a folder `f`
85    /// call: `t.fold_with(f)`.
86    ///
87    /// For most types, this just traverses the value, calling `fold_with`
88    /// on each field/element.
89    ///
90    /// For types of interest (such as `Ty`), the implementation of this method
91    /// calls a folder method specifically for that type (such as
92    /// `F::fold_ty`). This is where control transfers from `TypeFoldable`
93    /// to `TypeFolder`.
94    ///
95    /// Same as [`TypeFoldable::try_fold_with`], but not fallible. Make sure to keep
96    /// the behavior in sync across functions.
97    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self;
98}
99
100// This trait is implemented for types of interest.
101pub trait TypeSuperFoldable<I: Interner>: TypeFoldable<I> {
102    /// Provides a default fold for a recursive type of interest. This should
103    /// only be called within `TypeFolder` methods, when a non-custom traversal
104    /// is desired for the value of the type of interest passed to that method.
105    /// For example, in `MyFolder::try_fold_ty(ty)`, it is valid to call
106    /// `ty.try_super_fold_with(self)`, but any other folding should be done
107    /// with `xyz.try_fold_with(self)`.
108    fn try_super_fold_with<F: FallibleTypeFolder<I>>(
109        self,
110        folder: &mut F,
111    ) -> Result<Self, F::Error>;
112
113    /// A convenient alternative to `try_super_fold_with` for use with
114    /// infallible folders. Do not override this method, to ensure coherence
115    /// with `try_super_fold_with`.
116    fn super_fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self;
117}
118
119/// This trait is implemented for every infallible folding traversal. There is
120/// a fold method defined for every type of interest. Each such method has a
121/// default that does an "identity" fold. Implementations of these methods
122/// often fall back to a `super_fold_with` method if the primary argument
123/// doesn't satisfy a particular condition.
124///
125/// A blanket implementation of [`FallibleTypeFolder`] will defer to
126/// the infallible methods of this trait to ensure that the two APIs
127/// are coherent.
128pub trait TypeFolder<I: Interner>: Sized {
129    fn cx(&self) -> I;
130
131    fn fold_binder<T>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T>
132    where
133        T: TypeFoldable<I>,
134    {
135        t.super_fold_with(self)
136    }
137
138    fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
139        t.super_fold_with(self)
140    }
141
142    // The default region folder is a no-op because `Region` is non-recursive
143    // and has no `super_fold_with` method to call.
144    fn fold_region(&mut self, r: I::Region) -> I::Region {
145        r
146    }
147
148    fn fold_const(&mut self, c: I::Const) -> I::Const {
149        c.super_fold_with(self)
150    }
151
152    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
153        p.super_fold_with(self)
154    }
155}
156
157/// This trait is implemented for every folding traversal. There is a fold
158/// method defined for every type of interest. Each such method has a default
159/// that does an "identity" fold.
160///
161/// A blanket implementation of this trait (that defers to the relevant
162/// method of [`TypeFolder`]) is provided for all infallible folders in
163/// order to ensure the two APIs are coherent.
164pub trait FallibleTypeFolder<I: Interner>: Sized {
165    type Error;
166
167    fn cx(&self) -> I;
168
169    fn try_fold_binder<T>(&mut self, t: ty::Binder<I, T>) -> Result<ty::Binder<I, T>, Self::Error>
170    where
171        T: TypeFoldable<I>,
172    {
173        t.try_super_fold_with(self)
174    }
175
176    fn try_fold_ty(&mut self, t: I::Ty) -> Result<I::Ty, Self::Error> {
177        t.try_super_fold_with(self)
178    }
179
180    // The default region folder is a no-op because `Region` is non-recursive
181    // and has no `super_fold_with` method to call.
182    fn try_fold_region(&mut self, r: I::Region) -> Result<I::Region, Self::Error> {
183        Ok(r)
184    }
185
186    fn try_fold_const(&mut self, c: I::Const) -> Result<I::Const, Self::Error> {
187        c.try_super_fold_with(self)
188    }
189
190    fn try_fold_predicate(&mut self, p: I::Predicate) -> Result<I::Predicate, Self::Error> {
191        p.try_super_fold_with(self)
192    }
193}
194
195///////////////////////////////////////////////////////////////////////////
196// Traversal implementations.
197
198impl<I: Interner, T: TypeFoldable<I>, U: TypeFoldable<I>> TypeFoldable<I> for (T, U) {
199    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<(T, U), F::Error> {
200        Ok((self.0.try_fold_with(folder)?, self.1.try_fold_with(folder)?))
201    }
202
203    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
204        (self.0.fold_with(folder), self.1.fold_with(folder))
205    }
206}
207
208impl<I: Interner, A: TypeFoldable<I>, B: TypeFoldable<I>, C: TypeFoldable<I>> TypeFoldable<I>
209    for (A, B, C)
210{
211    fn try_fold_with<F: FallibleTypeFolder<I>>(
212        self,
213        folder: &mut F,
214    ) -> Result<(A, B, C), F::Error> {
215        Ok((
216            self.0.try_fold_with(folder)?,
217            self.1.try_fold_with(folder)?,
218            self.2.try_fold_with(folder)?,
219        ))
220    }
221
222    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
223        (self.0.fold_with(folder), self.1.fold_with(folder), self.2.fold_with(folder))
224    }
225}
226
227impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Option<T> {
228    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
229        Ok(match self {
230            Some(v) => Some(v.try_fold_with(folder)?),
231            None => None,
232        })
233    }
234
235    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
236        Some(self?.fold_with(folder))
237    }
238}
239
240impl<I: Interner, T: TypeFoldable<I>, E: TypeFoldable<I>> TypeFoldable<I> for Result<T, E> {
241    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
242        Ok(match self {
243            Ok(v) => Ok(v.try_fold_with(folder)?),
244            Err(e) => Err(e.try_fold_with(folder)?),
245        })
246    }
247
248    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
249        match self {
250            Ok(v) => Ok(v.fold_with(folder)),
251            Err(e) => Err(e.fold_with(folder)),
252        }
253    }
254}
255
256fn fold_arc<T: Clone, E>(
257    mut arc: Arc<T>,
258    fold: impl FnOnce(T) -> Result<T, E>,
259) -> Result<Arc<T>, E> {
260    // We merely want to replace the contained `T`, if at all possible,
261    // so that we don't needlessly allocate a new `Arc` or indeed clone
262    // the contained type.
263    unsafe {
264        // First step is to ensure that we have a unique reference to
265        // the contained type, which `Arc::make_mut` will accomplish (by
266        // allocating a new `Arc` and cloning the `T` only if required).
267        // This is done *before* casting to `Arc<ManuallyDrop<T>>` so that
268        // panicking during `make_mut` does not leak the `T`.
269        Arc::make_mut(&mut arc);
270
271        // Casting to `Arc<ManuallyDrop<T>>` is safe because `ManuallyDrop`
272        // is `repr(transparent)`.
273        let ptr = Arc::into_raw(arc).cast::<mem::ManuallyDrop<T>>();
274        let mut unique = Arc::from_raw(ptr);
275
276        // Call to `Arc::make_mut` above guarantees that `unique` is the
277        // sole reference to the contained value, so we can avoid doing
278        // a checked `get_mut` here.
279        let slot = Arc::get_mut(&mut unique).unwrap_unchecked();
280
281        // Semantically move the contained type out from `unique`, fold
282        // it, then move the folded value back into `unique`. Should
283        // folding fail, `ManuallyDrop` ensures that the "moved-out"
284        // value is not re-dropped.
285        let owned = mem::ManuallyDrop::take(slot);
286        let folded = fold(owned)?;
287        *slot = mem::ManuallyDrop::new(folded);
288
289        // Cast back to `Arc<T>`.
290        Ok(Arc::from_raw(Arc::into_raw(unique).cast()))
291    }
292}
293
294impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Arc<T> {
295    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
296        fold_arc(self, |t| t.try_fold_with(folder))
297    }
298
299    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
300        match fold_arc::<T, Infallible>(self, |t| Ok(t.fold_with(folder))) {
301            Ok(t) => t,
302        }
303    }
304}
305
306impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Box<T> {
307    fn try_fold_with<F: FallibleTypeFolder<I>>(mut self, folder: &mut F) -> Result<Self, F::Error> {
308        *self = (*self).try_fold_with(folder)?;
309        Ok(self)
310    }
311
312    fn fold_with<F: TypeFolder<I>>(mut self, folder: &mut F) -> Self {
313        *self = (*self).fold_with(folder);
314        self
315    }
316}
317
318impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Vec<T> {
319    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
320        self.into_iter().map(|t| t.try_fold_with(folder)).collect()
321    }
322
323    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
324        self.into_iter().map(|t| t.fold_with(folder)).collect()
325    }
326}
327
328impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for ThinVec<T> {
329    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
330        self.into_iter().map(|t| t.try_fold_with(folder)).collect()
331    }
332
333    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
334        self.into_iter().map(|t| t.fold_with(folder)).collect()
335    }
336}
337
338impl<I: Interner, T: TypeFoldable<I>> TypeFoldable<I> for Box<[T]> {
339    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
340        Vec::from(self).try_fold_with(folder).map(Vec::into_boxed_slice)
341    }
342
343    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
344        Vec::into_boxed_slice(Vec::from(self).fold_with(folder))
345    }
346}
347
348impl<I: Interner, T: TypeFoldable<I>, Ix: Idx> TypeFoldable<I> for IndexVec<Ix, T> {
349    fn try_fold_with<F: FallibleTypeFolder<I>>(self, folder: &mut F) -> Result<Self, F::Error> {
350        self.raw.try_fold_with(folder).map(IndexVec::from_raw)
351    }
352
353    fn fold_with<F: TypeFolder<I>>(self, folder: &mut F) -> Self {
354        IndexVec::from_raw(self.raw.fold_with(folder))
355    }
356}
357
358///////////////////////////////////////////////////////////////////////////
359// Shifter
360//
361// Shifts the De Bruijn indices on all escaping bound vars by a
362// fixed amount. Useful in instantiation or when otherwise introducing
363// a binding level that is not intended to capture the existing bound
364// vars. See comment on `shift_vars_through_binders` method in
365// `rustc_middle/src/ty/generic_args.rs` for more details.
366
367struct Shifter<I: Interner> {
368    cx: I,
369    current_index: ty::DebruijnIndex,
370    amount: u32,
371}
372
373impl<I: Interner> Shifter<I> {
374    fn new(cx: I, amount: u32) -> Self {
375        Shifter { cx, current_index: ty::INNERMOST, amount }
376    }
377}
378
379impl<I: Interner> TypeFolder<I> for Shifter<I> {
380    fn cx(&self) -> I {
381        self.cx
382    }
383
384    fn fold_binder<T: TypeFoldable<I>>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T> {
385        self.current_index.shift_in(1);
386        let t = t.super_fold_with(self);
387        self.current_index.shift_out(1);
388        t
389    }
390
391    fn fold_region(&mut self, r: I::Region) -> I::Region {
392        match r.kind() {
393            ty::ReBound(debruijn, br) if debruijn >= self.current_index => {
394                let debruijn = debruijn.shifted_in(self.amount);
395                Region::new_bound(self.cx, debruijn, br)
396            }
397            _ => r,
398        }
399    }
400
401    fn fold_ty(&mut self, ty: I::Ty) -> I::Ty {
402        match ty.kind() {
403            ty::Bound(debruijn, bound_ty) if debruijn >= self.current_index => {
404                let debruijn = debruijn.shifted_in(self.amount);
405                Ty::new_bound(self.cx, debruijn, bound_ty)
406            }
407
408            _ if ty.has_vars_bound_at_or_above(self.current_index) => ty.super_fold_with(self),
409            _ => ty,
410        }
411    }
412
413    fn fold_const(&mut self, ct: I::Const) -> I::Const {
414        match ct.kind() {
415            ty::ConstKind::Bound(debruijn, bound_ct) if debruijn >= self.current_index => {
416                let debruijn = debruijn.shifted_in(self.amount);
417                Const::new_bound(self.cx, debruijn, bound_ct)
418            }
419            _ => ct.super_fold_with(self),
420        }
421    }
422
423    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
424        if p.has_vars_bound_at_or_above(self.current_index) { p.super_fold_with(self) } else { p }
425    }
426}
427
428pub fn shift_region<I: Interner>(cx: I, region: I::Region, amount: u32) -> I::Region {
429    match region.kind() {
430        ty::ReBound(debruijn, br) if amount > 0 => {
431            Region::new_bound(cx, debruijn.shifted_in(amount), br)
432        }
433        _ => region,
434    }
435}
436
437#[instrument(level = "trace", skip(cx), ret)]
438pub fn shift_vars<I: Interner, T>(cx: I, value: T, amount: u32) -> T
439where
440    T: TypeFoldable<I>,
441{
442    if amount == 0 || !value.has_escaping_bound_vars() {
443        value
444    } else {
445        value.fold_with(&mut Shifter::new(cx, amount))
446    }
447}
448
449///////////////////////////////////////////////////////////////////////////
450// Region folder
451
452pub fn fold_regions<I: Interner, T>(
453    cx: I,
454    value: T,
455    f: impl FnMut(I::Region, ty::DebruijnIndex) -> I::Region,
456) -> T
457where
458    T: TypeFoldable<I>,
459{
460    value.fold_with(&mut RegionFolder::new(cx, f))
461}
462
463/// Folds over the substructure of a type, visiting its component
464/// types and all regions that occur *free* within it.
465///
466/// That is, function pointer types and trait object can introduce
467/// new bound regions which are not visited by this visitors as
468/// they are not free; only regions that occur free will be
469/// visited by `fld_r`.
470pub struct RegionFolder<I, F> {
471    cx: I,
472
473    /// Stores the index of a binder *just outside* the stuff we have
474    /// visited. So this begins as INNERMOST; when we pass through a
475    /// binder, it is incremented (via `shift_in`).
476    current_index: ty::DebruijnIndex,
477
478    /// Callback invokes for each free region. The `DebruijnIndex`
479    /// points to the binder *just outside* the ones we have passed
480    /// through.
481    fold_region_fn: F,
482}
483
484impl<I, F> RegionFolder<I, F> {
485    #[inline]
486    pub fn new(cx: I, fold_region_fn: F) -> RegionFolder<I, F> {
487        RegionFolder { cx, current_index: ty::INNERMOST, fold_region_fn }
488    }
489}
490
491impl<I, F> TypeFolder<I> for RegionFolder<I, F>
492where
493    I: Interner,
494    F: FnMut(I::Region, ty::DebruijnIndex) -> I::Region,
495{
496    fn cx(&self) -> I {
497        self.cx
498    }
499
500    fn fold_binder<T: TypeFoldable<I>>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T> {
501        self.current_index.shift_in(1);
502        let t = t.super_fold_with(self);
503        self.current_index.shift_out(1);
504        t
505    }
506
507    #[instrument(skip(self), level = "debug", ret)]
508    fn fold_region(&mut self, r: I::Region) -> I::Region {
509        match r.kind() {
510            ty::ReBound(debruijn, _) if debruijn < self.current_index => {
511                debug!(?self.current_index, "skipped bound region");
512                r
513            }
514            _ => {
515                debug!(?self.current_index, "folding free region");
516                (self.fold_region_fn)(r, self.current_index)
517            }
518        }
519    }
520
521    fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
522        if t.has_type_flags(
523            TypeFlags::HAS_FREE_REGIONS | TypeFlags::HAS_RE_BOUND | TypeFlags::HAS_RE_ERASED,
524        ) {
525            t.super_fold_with(self)
526        } else {
527            t
528        }
529    }
530
531    fn fold_const(&mut self, ct: I::Const) -> I::Const {
532        if ct.has_type_flags(
533            TypeFlags::HAS_FREE_REGIONS | TypeFlags::HAS_RE_BOUND | TypeFlags::HAS_RE_ERASED,
534        ) {
535            ct.super_fold_with(self)
536        } else {
537            ct
538        }
539    }
540
541    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
542        if p.has_type_flags(
543            TypeFlags::HAS_FREE_REGIONS | TypeFlags::HAS_RE_BOUND | TypeFlags::HAS_RE_ERASED,
544        ) {
545            p.super_fold_with(self)
546        } else {
547            p
548        }
549    }
550}