Skip to main content

cl_ast/ast/
fold.rs

1//! A folder (implementer of the [`Fold`] trait) maps ASTs to ASTs
2#![warn(clippy::all, clippy::pedantic)]
3#![allow(clippy::wildcard_imports, clippy::missing_errors_doc)]
4
5use super::*;
6
7/// Deconstructs an entire AST, and reconstructs it from parts.
8///
9/// Each function acts as a customization point.
10///
11/// Aside from [`AstTypes`], each node in the AST implements the [`Foldable`] trait,
12/// which provides double dispatch.
13pub trait Fold<From: AstTypes, To: AstTypes = From> {
14    type Error;
15
16    /// Consumes an Annotation in A, possibly transforms it, and produces a replacement Annotation
17    /// in B
18    fn fold_annotation(&mut self, anno: From::Annotation) -> Result<To::Annotation, Self::Error>;
19
20    /// Consumes a `MacroId` in A, possibly transforms it, and produces a replacement `MacroId` in B
21    fn fold_macro_id(&mut self, name: From::MacroId) -> Result<To::MacroId, Self::Error>;
22
23    /// Consumes a `Symbol` in A, possibly transforms it, and produces a replacement `Symbol` in B
24    fn fold_symbol(&mut self, name: From::Symbol) -> Result<To::Symbol, Self::Error>;
25
26    /// Consumes a `Path` in A, possibly transforms it, and produces a replacement `Path` in B
27    fn fold_path(&mut self, path: From::Path) -> Result<To::Path, Self::Error>;
28
29    /// Consumes a `Literal` in A, possibly transforms it, and produces a replacement `Literal` in B
30    fn fold_literal(&mut self, lit: From::Literal) -> Result<To::Literal, Self::Error>;
31
32    /// Folds an annotated expression, so the expression and annotation can be seen at once.
33    fn fold_at_expr(
34        &mut self,
35        expr: At<Expr<From>, From>,
36    ) -> Result<At<Expr<To>, To>, Self::Error> {
37        expr.children(self)
38    }
39
40    /// Consumes an [`Expr`], possibly transforms it, and produces a replacement [`Expr`]
41    fn fold_expr(&mut self, expr: Expr<From>) -> Result<Expr<To>, Self::Error> {
42        expr.children(self)
43    }
44
45    /// Consumes a [`Use`], possibly transforms it, and produces a replacement [`Use`]
46    fn fold_use(&mut self, item: Use<From>) -> Result<Use<To>, Self::Error> {
47        item.children(self)
48    }
49
50    /// Consumes a [`Pat`], possibly transforms it, and produces a replacement [`Pat`]
51    fn fold_pat(&mut self, pat: Pat<From>) -> Result<Pat<To>, Self::Error> {
52        pat.children(self)
53    }
54
55    /// Consumes a [`Bind`], possibly transforms it, and produces a replacement [`Bind`]
56    fn fold_bind(&mut self, bind: Bind<From>) -> Result<Bind<To>, Self::Error> {
57        bind.children(self)
58    }
59
60    /// Consumes a [`Make`], possibly transforms it, and produces a replacement [`Make`]
61    fn fold_make(&mut self, make: Make<From>) -> Result<Make<To>, Self::Error> {
62        make.children(self)
63    }
64
65    /// Consumes a [`MakeArm`], possibly transforms it, and produces a replacement [`MakeArm`]
66    fn fold_makearm(&mut self, arm: MakeArm<From>) -> Result<MakeArm<To>, Self::Error> {
67        arm.children(self)
68    }
69
70    /// Consumes a [`Make`], possibly transforms it, and produces a replacement [`Make`]
71    fn fold_match(&mut self, mtch: Match<From>) -> Result<Match<To>, Self::Error> {
72        mtch.children(self)
73    }
74
75    /// Consumes a [`MakeArm`], possibly transforms it, and produces a replacement [`MakeArm`]
76    fn fold_matcharm(&mut self, arm: MatchArm<From>) -> Result<MatchArm<To>, Self::Error> {
77        arm.children(self)
78    }
79}
80
81pub trait Foldable<A: AstTypes, B: AstTypes>: Sized {
82    /// The return type of the associated [Fold] function
83    type Out;
84
85    /// Calls `Self`'s appropriate [Folder](Fold) function(s)
86    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error>;
87
88    /// Destructures `self`, calling [`Foldable::fold_in`] on all foldable members,
89    /// and rebuilds a `Self` out of the results.
90    #[allow(unused_variables)]
91    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error>;
92}
93
94impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Expr<A> {
95    type Out = Expr<B>;
96
97    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
98        folder.fold_expr(self)
99    }
100
101    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
102        Ok(match self {
103            Self::Omitted => Expr::Omitted,
104            Self::Id(path) => Expr::Id(folder.fold_path(path)?),
105            Self::MetId(id) => Expr::MetId(folder.fold_macro_id(id)?),
106            Self::Lit(lit) => Expr::Lit(folder.fold_literal(lit)?),
107            Self::Use(item) => Expr::Use(item.fold_in(folder)?),
108            Self::Bind(bind) => Expr::Bind(bind.fold_in(folder)?),
109            Self::Make(make) => Expr::Make(make.fold_in(folder)?),
110            Self::Match(mtch) => Expr::Match(mtch.fold_in(folder)?),
111            Self::Op(op, annos) => Expr::Op(op, annos.fold_in(folder)?),
112        })
113    }
114}
115
116impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Use<A> {
117    type Out = Use<B>;
118
119    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
120        folder.fold_use(self)
121    }
122
123    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
124        Ok(match self {
125            Self::Glob => Use::Glob,
126            Self::Name(name) => Use::Name(folder.fold_symbol(name)?),
127            Self::Alias(name, alias) => {
128                Use::Alias(folder.fold_symbol(name)?, folder.fold_symbol(alias)?)
129            }
130            Self::Path(name, rest) => Use::Path(folder.fold_symbol(name)?, rest.fold_in(folder)?),
131            Self::Tree(items) => Use::Tree(items.fold_in(folder)?),
132        })
133    }
134}
135
136impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Pat<A> {
137    type Out = Pat<B>;
138
139    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
140        folder.fold_pat(self)
141    }
142
143    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
144        Ok(match self {
145            Self::Ignore => Pat::Ignore,
146            Self::Never => Pat::Never,
147            Self::MetId(name) => Pat::MetId(folder.fold_macro_id(name)?),
148            Self::Name(name) => Pat::Name(folder.fold_symbol(name)?),
149            Self::Value(expr) => Pat::Value(expr.fold_in(folder)?),
150            Self::Op(op, pats) => Pat::Op(op, pats.fold_in(folder)?),
151        })
152    }
153}
154
155impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Bind<A> {
156    type Out = Bind<B>;
157
158    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
159        folder.fold_bind(self)
160    }
161
162    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
163        let Self(op, gens, pat, exprs) = self;
164        Ok(Bind(
165            op,
166            gens.into_iter()
167                .map(|g| folder.fold_path(g))
168                .collect::<Result<_, _>>()?,
169            pat.fold_in(folder)?,
170            exprs.fold_in(folder)?,
171        ))
172    }
173}
174
175impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Make<A> {
176    type Out = Make<B>;
177
178    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
179        folder.fold_make(self)
180    }
181
182    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
183        let Self(expr, arms) = self;
184        Ok(Make(expr.fold_in(folder)?, arms.fold_in(folder)?))
185    }
186}
187
188impl<A: AstTypes, B: AstTypes> Foldable<A, B> for MakeArm<A> {
189    type Out = MakeArm<B>;
190
191    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
192        folder.fold_makearm(self)
193    }
194
195    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
196        let Self(name, expr) = self;
197        Ok(MakeArm(folder.fold_symbol(name)?, expr.fold_in(folder)?))
198    }
199}
200
201impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Match<A> {
202    type Out = Match<B>;
203
204    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
205        folder.fold_match(self)
206    }
207
208    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
209        let Self(scrutinee, arms) = self;
210        Ok(Match(scrutinee.fold_in(folder)?, arms.fold_in(folder)?))
211    }
212}
213
214impl<A: AstTypes, B: AstTypes> Foldable<A, B> for MatchArm<A> {
215    type Out = MatchArm<B>;
216
217    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
218        folder.fold_matcharm(self)
219    }
220
221    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
222        let Self(pat, expr) = self;
223        Ok(MatchArm(pat.fold_in(folder)?, expr.fold_in(folder)?))
224    }
225}
226
227impl<A: AstTypes, B: AstTypes> Foldable<A, B> for At<Expr<A>, A> {
228    type Out = At<Expr<B>, B>;
229
230    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
231        folder.fold_at_expr(self)
232    }
233
234    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
235        let Self(expr, anno) = self;
236        Ok(At(expr.children(folder)?, folder.fold_annotation(anno)?))
237    }
238}
239
240//////////////////////////////////////////////
241//  GENERIC IMPLEMENTATIONS ON COLLECTIONS  //
242//////////////////////////////////////////////
243
244// Maps the value in the box across `f()` without deallocating
245// fn box_try_map<T, E>(boxed: Box<T>, f: impl FnOnce(T) -> Result<T, E>) -> Result<Box<T>, E> {
246//     // TODO: replace with Box::take when it stabilizes.
247//     let rawbox = Box::into_raw(boxed);
248
249//     // Safety: `rawbox` came from a Box, so it is aligned and initialized.
250//     //     To prevent further reuse and deallocate on failure, rawbox is
251//     //     shadowed by a Box<MaybeUninit<T>>.
252//     // Safety: MaybeUninit<T> has the same size and alignment as T.
253//     let (value, rawbox) = (unsafe { rawbox.read() }, unsafe {
254//         Box::from_raw(rawbox.cast::<MaybeUninit<T>>())
255//     });
256
257//     // rawbox is reinitialized with f(value)
258//     Ok(Box::write(rawbox, f(value)?))
259// }
260
261impl<T, A, B> Foldable<A, B> for Box<T>
262where
263    T: Foldable<A, B>,
264    A: AstTypes,
265    B: AstTypes,
266{
267    type Out = Box<T::Out>;
268
269    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
270        let value = *self;
271        Ok(Box::new(value.fold_in(folder)?))
272    }
273
274    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
275        let value = *self;
276        Ok(Box::new(value.children(folder)?))
277    }
278}
279
280impl<T: Foldable<A, B>, A: AstTypes, B: AstTypes> Foldable<A, B> for Vec<T> {
281    type Out = Vec<T::Out>;
282
283    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
284        self.children(folder)
285    }
286
287    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
288        // TODO: is this correct for generic data structures?
289        self.into_iter().map(|e| e.fold_in(folder)).collect()
290    }
291}
292
293impl<T: Foldable<A, B>, A: AstTypes, B: AstTypes> Foldable<A, B> for Option<T> {
294    type Out = Option<T::Out>;
295
296    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
297        Ok(match self {
298            Self::Some(value) => Some(value.fold_in(folder)?),
299            Self::None => None,
300        })
301    }
302
303    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
304        Ok(match self {
305            Self::Some(value) => Some(value.children(folder)?),
306            Self::None => None,
307        })
308    }
309}