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    /// Folds an annotated pattern, so the pattern and annotation can be seen at once.
51    fn fold_at_pat(&mut self, pat: At<Pat<From>, From>) -> Result<At<Pat<To>, To>, Self::Error> {
52        pat.children(self)
53    }
54
55    /// Consumes a [`Pat`], possibly transforms it, and produces a replacement [`Pat`]
56    fn fold_pat(&mut self, pat: Pat<From>) -> Result<Pat<To>, Self::Error> {
57        pat.children(self)
58    }
59
60    /// Consumes a [`Bind`], possibly transforms it, and produces a replacement [`Bind`]
61    fn fold_bind(&mut self, bind: Bind<From>) -> Result<Bind<To>, Self::Error> {
62        bind.children(self)
63    }
64
65    /// Consumes a [`Make`], possibly transforms it, and produces a replacement [`Make`]
66    fn fold_make(&mut self, make: Make<From>) -> Result<Make<To>, Self::Error> {
67        make.children(self)
68    }
69
70    /// Consumes a [`MakeArm`], possibly transforms it, and produces a replacement [`MakeArm`]
71    fn fold_makearm(&mut self, arm: MakeArm<From>) -> Result<MakeArm<To>, Self::Error> {
72        arm.children(self)
73    }
74
75    /// Consumes a [`Make`], possibly transforms it, and produces a replacement [`Make`]
76    fn fold_match(&mut self, mtch: Match<From>) -> Result<Match<To>, Self::Error> {
77        mtch.children(self)
78    }
79
80    /// Consumes a [`MakeArm`], possibly transforms it, and produces a replacement [`MakeArm`]
81    fn fold_matcharm(&mut self, arm: MatchArm<From>) -> Result<MatchArm<To>, Self::Error> {
82        arm.children(self)
83    }
84}
85
86/// ```ignore
87/// pub macro impl_default_fold($Src: ty, $Dst: ty)
88/// ```
89///  Implements [`Into`]-based defaults for the required [`Fold`] members:
90/// - [`Fold::fold_annotation`] `where Src::Annotation: Into<Dst::Annotation>`
91/// - [`Fold::fold_macro_id`] `where Src::MacroId: Into<Dst::MacroId>`
92/// - [`Fold::fold_symbol`] `where Src::Symbol: Into<Dst::Symbol>`
93/// - [`Fold::fold_path`] `where Src::Path: Into<Dst::Path>`
94/// - [`Fold::fold_literal`] `where Src::Literal: Into<Dst::Literal>`
95///
96/// # Examples:
97/// Implements an "identity" folder
98/// ```rust
99/// # use cl_ast::ast::{AstTypes, Annotation};
100/// # use cl_ast::ast::fold::{Fold, impl_default_fold};
101/// struct IdentityFold;
102/// impl<A: AstTypes> Fold<A, A> for IdentityFold {
103///     type Error = std::convert::Infallible;
104///     impl_default_fold!(A, A);
105/// }
106/// ```
107pub macro impl_default_fold($Src: ty, $Dst: ty) {
108    fn fold_annotation(
109        &mut self,
110        anno: <$Src as AstTypes>::Annotation,
111    ) -> Result<<$Dst as AstTypes>::Annotation, Self::Error> {
112        Ok(anno.into())
113    }
114    fn fold_macro_id(
115        &mut self,
116        name: <$Src as AstTypes>::MacroId,
117    ) -> Result<<$Dst as AstTypes>::MacroId, Self::Error> {
118        Ok(name.into())
119    }
120    fn fold_symbol(
121        &mut self,
122        name: <$Src as AstTypes>::Symbol,
123    ) -> Result<<$Dst as AstTypes>::Symbol, Self::Error> {
124        Ok(name.into())
125    }
126    fn fold_path(
127        &mut self,
128        path: <$Src as AstTypes>::Path,
129    ) -> Result<<$Dst as AstTypes>::Path, Self::Error> {
130        Ok(path.into())
131    }
132    fn fold_literal(
133        &mut self,
134        lit: <$Src as AstTypes>::Literal,
135    ) -> Result<<$Dst as AstTypes>::Literal, Self::Error> {
136        Ok(lit.into())
137    }
138}
139
140/// Implements depth-first traversal for folders
141pub trait Foldable<A: AstTypes, B: AstTypes>: Sized {
142    /// The return type of the associated [Fold] function
143    type Out;
144
145    /// Calls `Self`'s appropriate [Folder](Fold) function(s)
146    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error>;
147
148    /// Destructures `self`, calling [`Foldable::fold_in`] on all foldable members,
149    /// and rebuilds a `Self` out of the results.
150    #[allow(unused_variables)]
151    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error>;
152}
153
154impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Expr<A> {
155    type Out = Expr<B>;
156
157    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
158        folder.fold_expr(self)
159    }
160
161    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
162        Ok(match self {
163            Self::Omitted => Expr::Omitted,
164            Self::Id(path) => Expr::Id(folder.fold_path(path)?),
165            Self::MetId(id) => Expr::MetId(folder.fold_macro_id(id)?),
166            Self::Lit(lit) => Expr::Lit(folder.fold_literal(lit)?),
167            Self::Use(item) => Expr::Use(item.fold_in(folder)?),
168            Self::Bind(bind) => Expr::Bind(bind.fold_in(folder)?),
169            Self::Make(make) => Expr::Make(make.fold_in(folder)?),
170            Self::Match(mtch) => Expr::Match(mtch.fold_in(folder)?),
171            Self::Op(op, annos) => Expr::Op(op, annos.fold_in(folder)?),
172        })
173    }
174}
175
176impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Use<A> {
177    type Out = Use<B>;
178
179    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
180        folder.fold_use(self)
181    }
182
183    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
184        Ok(match self {
185            Self::Glob => Use::Glob,
186            Self::Name(name) => Use::Name(folder.fold_symbol(name)?),
187            Self::Alias(name, alias) => {
188                Use::Alias(folder.fold_symbol(name)?, folder.fold_symbol(alias)?)
189            }
190            Self::Path(name, rest) => Use::Path(folder.fold_symbol(name)?, rest.fold_in(folder)?),
191            Self::Tree(items) => Use::Tree(items.fold_in(folder)?),
192        })
193    }
194}
195
196impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Pat<A> {
197    type Out = Pat<B>;
198
199    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
200        folder.fold_pat(self)
201    }
202
203    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
204        Ok(match self {
205            Self::Ignore => Pat::Ignore,
206            Self::Never => Pat::Never,
207            Self::MetId(name) => Pat::MetId(folder.fold_macro_id(name)?),
208            Self::Name(name) => Pat::Name(folder.fold_symbol(name)?),
209            Self::Value(expr) => Pat::Value(expr.fold_in(folder)?),
210            Self::Op(op, pats) => Pat::Op(op, pats.fold_in(folder)?),
211        })
212    }
213}
214
215impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Bind<A> {
216    type Out = Bind<B>;
217
218    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
219        folder.fold_bind(self)
220    }
221
222    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
223        let Self(op, gens, pat, exprs) = self;
224        Ok(Bind(
225            op,
226            gens.into_iter()
227                .map(|g| folder.fold_path(g))
228                .collect::<Result<_, _>>()?,
229            pat.fold_in(folder)?,
230            exprs.fold_in(folder)?,
231        ))
232    }
233}
234
235impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Make<A> {
236    type Out = Make<B>;
237
238    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
239        folder.fold_make(self)
240    }
241
242    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
243        let Self(expr, arms) = self;
244        Ok(Make(expr.fold_in(folder)?, arms.fold_in(folder)?))
245    }
246}
247
248impl<A: AstTypes, B: AstTypes> Foldable<A, B> for MakeArm<A> {
249    type Out = MakeArm<B>;
250
251    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
252        folder.fold_makearm(self)
253    }
254
255    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
256        let Self(name, expr) = self;
257        Ok(MakeArm(folder.fold_symbol(name)?, expr.fold_in(folder)?))
258    }
259}
260
261impl<A: AstTypes, B: AstTypes> Foldable<A, B> for Match<A> {
262    type Out = Match<B>;
263
264    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
265        folder.fold_match(self)
266    }
267
268    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
269        let Self(scrutinee, arms) = self;
270        Ok(Match(scrutinee.fold_in(folder)?, arms.fold_in(folder)?))
271    }
272}
273
274impl<A: AstTypes, B: AstTypes> Foldable<A, B> for MatchArm<A> {
275    type Out = MatchArm<B>;
276
277    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
278        folder.fold_matcharm(self)
279    }
280
281    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
282        let Self(pat, expr) = self;
283        Ok(MatchArm(pat.fold_in(folder)?, expr.fold_in(folder)?))
284    }
285}
286
287impl<A: AstTypes, B: AstTypes> Foldable<A, B> for At<Expr<A>, A> {
288    type Out = At<Expr<B>, B>;
289
290    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
291        folder.fold_at_expr(self)
292    }
293
294    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
295        let Self(expr, anno) = self;
296        Ok(At(expr.children(folder)?, folder.fold_annotation(anno)?))
297    }
298}
299
300impl<A: AstTypes, B: AstTypes> Foldable<A, B> for At<Pat<A>, A> {
301    type Out = At<Pat<B>, B>;
302
303    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
304        folder.fold_at_pat(self)
305    }
306
307    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
308        let Self(pat, anno) = self;
309        Ok(At(pat.children(folder)?, folder.fold_annotation(anno)?))
310    }
311}
312
313//////////////////////////////////////////////
314//  GENERIC IMPLEMENTATIONS ON COLLECTIONS  //
315//////////////////////////////////////////////
316
317// Maps the value in the box across `f()` without deallocating
318// fn box_try_map<T, E>(boxed: Box<T>, f: impl FnOnce(T) -> Result<T, E>) -> Result<Box<T>, E> {
319//     // TODO: replace with Box::take when it stabilizes.
320//     let rawbox = Box::into_raw(boxed);
321
322//     // Safety: `rawbox` came from a Box, so it is aligned and initialized.
323//     //     To prevent further reuse and deallocate on failure, rawbox is
324//     //     shadowed by a Box<MaybeUninit<T>>.
325//     // Safety: MaybeUninit<T> has the same size and alignment as T.
326//     let (value, rawbox) = (unsafe { rawbox.read() }, unsafe {
327//         Box::from_raw(rawbox.cast::<MaybeUninit<T>>())
328//     });
329
330//     // rawbox is reinitialized with f(value)
331//     Ok(Box::write(rawbox, f(value)?))
332// }
333
334impl<T, A, B> Foldable<A, B> for Box<T>
335where
336    T: Foldable<A, B>,
337    A: AstTypes,
338    B: AstTypes,
339{
340    type Out = Box<T::Out>;
341
342    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
343        let value = *self;
344        Ok(Box::new(value.fold_in(folder)?))
345    }
346
347    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
348        let value = *self;
349        Ok(Box::new(value.children(folder)?))
350    }
351}
352
353impl<T: Foldable<A, B>, A: AstTypes, B: AstTypes> Foldable<A, B> for Vec<T> {
354    type Out = Vec<T::Out>;
355
356    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
357        self.children(folder)
358    }
359
360    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
361        // TODO: is this correct for generic data structures?
362        self.into_iter().map(|e| e.fold_in(folder)).collect()
363    }
364}
365
366impl<T: Foldable<A, B>, A: AstTypes, B: AstTypes> Foldable<A, B> for Option<T> {
367    type Out = Option<T::Out>;
368
369    fn fold_in<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
370        Ok(match self {
371            Self::Some(value) => Some(value.fold_in(folder)?),
372            Self::None => None,
373        })
374    }
375
376    fn children<F: Fold<A, B> + ?Sized>(self, folder: &mut F) -> Result<Self::Out, F::Error> {
377        Ok(match self {
378            Self::Some(value) => Some(value.children(folder)?),
379            Self::None => None,
380        })
381    }
382}