blueprint_starlark_syntax/syntax/
payload_map.rs

1/*
2 * Copyright 2018 The Starlark in Rust Authors.
3 * Copyright (c) Facebook, Inc. and its affiliates.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     https://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 */
17
18//! Map AST payload.
19
20use crate::codemap::Spanned;
21use crate::slice_vec_ext::VecExt;
22use crate::syntax::ast::ArgumentP;
23use crate::syntax::ast::AssignIdentP;
24use crate::syntax::ast::AssignP;
25use crate::syntax::ast::AssignTargetP;
26use crate::syntax::ast::AstPayload;
27use crate::syntax::ast::CallArgsP;
28use crate::syntax::ast::ClauseP;
29use crate::syntax::ast::DefP;
30use crate::syntax::ast::ExprP;
31use crate::syntax::ast::FStringP;
32use crate::syntax::ast::ForClauseP;
33use crate::syntax::ast::ForP;
34use crate::syntax::ast::IdentP;
35use crate::syntax::ast::LambdaP;
36use crate::syntax::ast::LoadArgP;
37use crate::syntax::ast::LoadP;
38use crate::syntax::ast::MatchP;
39use crate::syntax::ast::CaseClauseP;
40use crate::syntax::ast::ParameterP;
41use crate::syntax::ast::StmtP;
42use crate::syntax::ast::StructFieldP;
43use crate::syntax::ast::StructP;
44use crate::syntax::ast::TypeExprP;
45
46pub trait AstPayloadFunction<A: AstPayload, B: AstPayload> {
47    fn map_load(&mut self, import_path: &str, a: A::LoadPayload) -> B::LoadPayload;
48    fn map_ident(&mut self, a: A::IdentPayload) -> B::IdentPayload;
49    fn map_ident_assign(&mut self, a: A::IdentAssignPayload) -> B::IdentAssignPayload;
50    fn map_def(&mut self, a: A::DefPayload) -> B::DefPayload;
51    fn map_type_expr(&mut self, a: A::TypeExprPayload) -> B::TypeExprPayload;
52}
53
54impl<A: AstPayload> LoadArgP<A> {
55    pub fn into_map_payload<B: AstPayload>(
56        self,
57        f: &mut impl AstPayloadFunction<A, B>,
58    ) -> LoadArgP<B> {
59        let LoadArgP {
60            local,
61            their: remote,
62            comma,
63        } = self;
64        LoadArgP {
65            local: local.into_map_payload(f),
66            their: remote,
67            comma,
68        }
69    }
70}
71
72impl<A: AstPayload> LoadP<A> {
73    pub fn into_map_payload<B: AstPayload>(
74        self,
75        f: &mut impl AstPayloadFunction<A, B>,
76    ) -> LoadP<B> {
77        let LoadP {
78            module,
79            args,
80            payload,
81        } = self;
82        let payload = f.map_load(&module.node, payload);
83        LoadP {
84            module,
85            args: args.into_map(|a| a.into_map_payload(f)),
86            payload,
87        }
88    }
89}
90
91impl<A: AstPayload> AssignP<A> {
92    pub fn into_map_payload<B: AstPayload>(
93        self,
94        f: &mut impl AstPayloadFunction<A, B>,
95    ) -> AssignP<B> {
96        let AssignP { lhs, ty, rhs } = self;
97        AssignP {
98            lhs: lhs.into_map_payload(f),
99            ty: ty.map(|ty| ty.into_map_payload(f)),
100            rhs: rhs.into_map_payload(f),
101        }
102    }
103}
104
105impl<A: AstPayload> ForP<A> {
106    pub fn into_map_payload<B: AstPayload>(self, f: &mut impl AstPayloadFunction<A, B>) -> ForP<B> {
107        let ForP { var, over, body } = self;
108        ForP {
109            var: var.into_map_payload(f),
110            over: over.into_map_payload(f),
111            body: Box::new(body.into_map_payload(f)),
112        }
113    }
114}
115
116impl<A: AstPayload> StmtP<A> {
117    pub fn into_map_payload<B: AstPayload>(
118        self,
119        f: &mut impl AstPayloadFunction<A, B>,
120    ) -> StmtP<B> {
121        match self {
122            StmtP::Break => StmtP::Break,
123            StmtP::Continue => StmtP::Continue,
124            StmtP::Pass => StmtP::Pass,
125            StmtP::Return(None) => StmtP::Return(None),
126            StmtP::Return(Some(e)) => StmtP::Return(Some(e.into_map_payload(f))),
127            StmtP::Yield(None) => StmtP::Yield(None),
128            StmtP::Yield(Some(e)) => StmtP::Yield(Some(e.into_map_payload(f))),
129            StmtP::Expression(e) => StmtP::Expression(e.into_map_payload(f)),
130            StmtP::Assign(assign) => StmtP::Assign(assign.into_map_payload(f)),
131            StmtP::AssignModify(lhs, op, rhs) => StmtP::AssignModify(
132                lhs.into_map_payload(f),
133                op,
134                Box::new(rhs.into_map_payload(f)),
135            ),
136            StmtP::Statements(stmts) => {
137                StmtP::Statements(stmts.into_map(|s| s.into_map_payload(f)))
138            }
139            StmtP::If(cond, then_block) => StmtP::If(
140                cond.into_map_payload(f),
141                Box::new(then_block.into_map_payload(f)),
142            ),
143            StmtP::IfElse(cond, then_block_else_block) => {
144                let (then_block, else_block) = *then_block_else_block;
145                StmtP::IfElse(
146                    cond.into_map_payload(f),
147                    Box::new((
148                        then_block.into_map_payload(f),
149                        else_block.into_map_payload(f),
150                    )),
151                )
152            }
153            StmtP::For(fr) => StmtP::For(fr.into_map_payload(f)),
154            StmtP::Def(DefP {
155                name,
156                params,
157                return_type,
158                body,
159                payload,
160            }) => StmtP::Def(DefP {
161                name: name.into_map_payload(f),
162                params: params.into_map(|p| p.into_map_payload(f)),
163                return_type: return_type.map(|ret| Box::new(ret.into_map_payload(f))),
164                body: Box::new(body.into_map_payload(f)),
165                payload: f.map_def(payload),
166            }),
167            StmtP::Load(load) => StmtP::Load(load.into_map_payload(f)),
168            StmtP::Struct(s) => StmtP::Struct(s.into_map_payload(f)),
169            StmtP::Match(m) => StmtP::Match(m.into_map_payload(f)),
170        }
171    }
172}
173
174impl<A: AstPayload> StructP<A> {
175    pub fn into_map_payload<B: AstPayload>(
176        self,
177        f: &mut impl AstPayloadFunction<A, B>,
178    ) -> StructP<B> {
179        let StructP { name, fields } = self;
180        StructP {
181            name: name.into_map_payload(f),
182            fields: fields.into_map(|field| field.into_map_payload(f)),
183        }
184    }
185}
186
187impl<A: AstPayload> StructFieldP<A> {
188    pub fn into_map_payload<B: AstPayload>(
189        self,
190        f: &mut impl AstPayloadFunction<A, B>,
191    ) -> StructFieldP<B> {
192        let StructFieldP { name, typ, default } = self;
193        StructFieldP {
194            name: name.into_map_payload(f),
195            typ: typ.into_map_payload(f),
196            default: default.map(|e| e.into_map_payload(f)),
197        }
198    }
199}
200
201impl<A: AstPayload> MatchP<A> {
202    pub fn into_map_payload<B: AstPayload>(
203        self,
204        f: &mut impl AstPayloadFunction<A, B>,
205    ) -> MatchP<B> {
206        let MatchP { subject, cases } = self;
207        MatchP {
208            subject: subject.into_map_payload(f),
209            cases: cases.into_map(|c| c.map(|cc| cc.into_map_payload(f))),
210        }
211    }
212}
213
214impl<A: AstPayload> CaseClauseP<A> {
215    pub fn into_map_payload<B: AstPayload>(
216        self,
217        f: &mut impl AstPayloadFunction<A, B>,
218    ) -> CaseClauseP<B> {
219        let CaseClauseP { pattern, guard, body } = self;
220        CaseClauseP {
221            pattern: pattern.into_map_payload(f),
222            guard: guard.map(|g| g.into_map_payload(f)),
223            body: body.into_map_payload(f),
224        }
225    }
226}
227
228impl<A: AstPayload> ExprP<A> {
229    pub fn into_map_payload<B: AstPayload>(
230        self,
231        f: &mut impl AstPayloadFunction<A, B>,
232    ) -> ExprP<B> {
233        match self {
234            ExprP::Tuple(exprs) => ExprP::Tuple(exprs.into_map(|e| e.into_map_payload(f))),
235            ExprP::Dot(object, field) => ExprP::Dot(Box::new(object.into_map_payload(f)), field),
236            ExprP::Call(ca, args) => ExprP::Call(
237                Box::new(ca.into_map_payload(f)),
238                CallArgsP {
239                    args: args.args.into_map(|a| a.into_map_payload(f)),
240                },
241            ),
242            ExprP::Index(array_index) => {
243                let (array, index) = *array_index;
244                ExprP::Index(Box::new((
245                    array.into_map_payload(f),
246                    index.into_map_payload(f),
247                )))
248            }
249            ExprP::Index2(a_i0_i1) => {
250                let (array, i0, i1) = *a_i0_i1;
251                ExprP::Index2(Box::new((
252                    array.into_map_payload(f),
253                    i0.into_map_payload(f),
254                    i1.into_map_payload(f),
255                )))
256            }
257            ExprP::Slice(x, a, b, c) => ExprP::Slice(
258                Box::new(x.into_map_payload(f)),
259                a.map(|e| Box::new(e.into_map_payload(f))),
260                b.map(|e| Box::new(e.into_map_payload(f))),
261                c.map(|e| Box::new(e.into_map_payload(f))),
262            ),
263            ExprP::Identifier(id) => ExprP::Identifier(id.into_map_payload(f)),
264            ExprP::Lambda(LambdaP {
265                params,
266                body,
267                payload,
268            }) => ExprP::Lambda(LambdaP {
269                params: params.into_map(|p| p.into_map_payload(f)),
270                body: Box::new(body.into_map_payload(f)),
271                payload: f.map_def(payload),
272            }),
273            ExprP::Literal(l) => ExprP::Literal(l),
274            ExprP::Not(e) => ExprP::Not(Box::new(e.into_map_payload(f))),
275            ExprP::Minus(e) => ExprP::Minus(Box::new(e.into_map_payload(f))),
276            ExprP::Plus(e) => ExprP::Plus(Box::new(e.into_map_payload(f))),
277            ExprP::BitNot(e) => ExprP::BitNot(Box::new(e.into_map_payload(f))),
278            ExprP::Op(l, op, r) => ExprP::Op(
279                Box::new(l.into_map_payload(f)),
280                op,
281                Box::new(r.into_map_payload(f)),
282            ),
283            ExprP::If(a_b_c) => {
284                let (a, b, c) = *a_b_c;
285                ExprP::If(Box::new((
286                    a.into_map_payload(f),
287                    b.into_map_payload(f),
288                    c.into_map_payload(f),
289                )))
290            }
291            ExprP::List(es) => ExprP::List(es.into_map(|e| e.into_map_payload(f))),
292            ExprP::Dict(kvs) => {
293                ExprP::Dict(kvs.into_map(|(k, v)| (k.into_map_payload(f), v.into_map_payload(f))))
294            }
295            ExprP::Set(es) => ExprP::Set(es.into_map(|e| e.into_map_payload(f))),
296            ExprP::ListComprehension(e, c0, cs) => ExprP::ListComprehension(
297                Box::new(e.into_map_payload(f)),
298                Box::new(c0.into_map_payload(f)),
299                cs.into_map(|c| c.into_map_payload(f)),
300            ),
301            ExprP::SetComprehension(e, c0, cs) => ExprP::SetComprehension(
302                Box::new(e.into_map_payload(f)),
303                Box::new(c0.into_map_payload(f)),
304                cs.into_map(|c| c.into_map_payload(f)),
305            ),
306            ExprP::DictComprehension(k_v, c0, cs) => {
307                let (k, v) = *k_v;
308                ExprP::DictComprehension(
309                    Box::new((k.into_map_payload(f), v.into_map_payload(f))),
310                    Box::new(c0.into_map_payload(f)),
311                    cs.into_map(|c| c.into_map_payload(f)),
312                )
313            }
314            ExprP::FString(fstring) => ExprP::FString(fstring.into_map_payload(f)),
315        }
316    }
317}
318
319impl<A: AstPayload> TypeExprP<A> {
320    pub fn into_map_payload<B: AstPayload>(
321        self,
322        f: &mut impl AstPayloadFunction<A, B>,
323    ) -> TypeExprP<B> {
324        let TypeExprP { expr, payload } = self;
325        TypeExprP {
326            expr: expr.map(|e| e.into_map_payload(f)),
327            payload: f.map_type_expr(payload),
328        }
329    }
330}
331
332impl<A: AstPayload> AssignTargetP<A> {
333    pub fn into_map_payload<B: AstPayload>(
334        self,
335        f: &mut impl AstPayloadFunction<A, B>,
336    ) -> AssignTargetP<B> {
337        match self {
338            AssignTargetP::Tuple(args) => {
339                AssignTargetP::Tuple(args.into_map(|a| a.into_map_payload(f)))
340            }
341            AssignTargetP::Index(array_index) => {
342                let (array, index) = *array_index;
343                AssignTargetP::Index(Box::new((
344                    array.into_map_payload(f),
345                    index.into_map_payload(f),
346                )))
347            }
348            AssignTargetP::Dot(object, field) => {
349                AssignTargetP::Dot(Box::new(object.into_map_payload(f)), field)
350            }
351            AssignTargetP::Identifier(ident) => {
352                AssignTargetP::Identifier(ident.into_map_payload(f))
353            }
354        }
355    }
356}
357
358impl<A: AstPayload> AssignIdentP<A> {
359    pub fn into_map_payload<B: AstPayload>(
360        self,
361        f: &mut impl AstPayloadFunction<A, B>,
362    ) -> AssignIdentP<B> {
363        let AssignIdentP {
364            ident: s,
365            payload: p,
366        } = self;
367        AssignIdentP {
368            ident: s,
369            payload: f.map_ident_assign(p),
370        }
371    }
372}
373
374impl<A: AstPayload> IdentP<A> {
375    pub fn into_map_payload<B: AstPayload>(
376        self,
377        f: &mut impl AstPayloadFunction<A, B>,
378    ) -> IdentP<B> {
379        let IdentP {
380            ident: s,
381            payload: p,
382        } = self;
383        IdentP {
384            ident: s,
385            payload: f.map_ident(p),
386        }
387    }
388}
389
390impl<A: AstPayload> ParameterP<A> {
391    pub fn into_map_payload<B: AstPayload>(
392        self,
393        f: &mut impl AstPayloadFunction<A, B>,
394    ) -> ParameterP<B> {
395        match self {
396            ParameterP::Normal(name, ty, defa) => ParameterP::Normal(
397                name.into_map_payload(f),
398                ty.map(|defa| Box::new(defa.into_map_payload(f))),
399                defa.map(|defa| Box::new(defa.into_map_payload(f))),
400            ),
401            ParameterP::NoArgs => ParameterP::NoArgs,
402            ParameterP::Slash => ParameterP::Slash,
403            ParameterP::Args(name, ty) => ParameterP::Args(
404                name.into_map_payload(f),
405                ty.map(|defa| Box::new(defa.into_map_payload(f))),
406            ),
407            ParameterP::KwArgs(name, ty) => ParameterP::KwArgs(
408                name.into_map_payload(f),
409                ty.map(|defa| Box::new(defa.into_map_payload(f))),
410            ),
411        }
412    }
413}
414
415impl<A: AstPayload> ArgumentP<A> {
416    pub fn into_map_payload<B: AstPayload>(
417        self,
418        f: &mut impl AstPayloadFunction<A, B>,
419    ) -> ArgumentP<B> {
420        match self {
421            ArgumentP::Positional(e) => ArgumentP::Positional(e.into_map_payload(f)),
422            ArgumentP::Named(n, e) => ArgumentP::Named(n, e.into_map_payload(f)),
423            ArgumentP::Args(e) => ArgumentP::Args(e.into_map_payload(f)),
424            ArgumentP::KwArgs(e) => ArgumentP::KwArgs(e.into_map_payload(f)),
425        }
426    }
427}
428
429impl<A: AstPayload> ClauseP<A> {
430    pub fn into_map_payload<B: AstPayload>(
431        self,
432        f: &mut impl AstPayloadFunction<A, B>,
433    ) -> ClauseP<B> {
434        match self {
435            ClauseP::For(c) => ClauseP::For(c.into_map_payload(f)),
436            ClauseP::If(e) => ClauseP::If(e.into_map_payload(f)),
437        }
438    }
439}
440
441impl<A: AstPayload> ForClauseP<A> {
442    pub fn into_map_payload<B: AstPayload>(
443        self,
444        f: &mut impl AstPayloadFunction<A, B>,
445    ) -> ForClauseP<B> {
446        ForClauseP {
447            var: self.var.into_map_payload(f),
448            over: self.over.into_map_payload(f),
449        }
450    }
451}
452
453impl<A: AstPayload> FStringP<A> {
454    pub fn into_map_payload<B: AstPayload>(
455        self,
456        f: &mut impl AstPayloadFunction<A, B>,
457    ) -> FStringP<B> {
458        let Self {
459            format,
460            expressions,
461        } = self;
462
463        FStringP {
464            format,
465            expressions: expressions.into_map(|id| id.into_map_payload(f)),
466        }
467    }
468}
469
470macro_rules! ast_payload_map_stub {
471    ($ast_type:ident, $ext:ident) => {
472        pub trait $ext<A: AstPayload> {
473            fn into_map_payload<B: AstPayload>(
474                self,
475                f: &mut impl AstPayloadFunction<A, B>,
476            ) -> Spanned<$ast_type<B>>;
477        }
478
479        impl<A: AstPayload> $ext<A> for Spanned<$ast_type<A>> {
480            fn into_map_payload<B: AstPayload>(
481                self,
482                f: &mut impl AstPayloadFunction<A, B>,
483            ) -> Spanned<$ast_type<B>> {
484                let Spanned { span, node } = self;
485                Spanned {
486                    span,
487                    node: node.into_map_payload(f),
488                }
489            }
490        }
491    };
492}
493
494ast_payload_map_stub!(ExprP, ExprPExt);
495ast_payload_map_stub!(TypeExprP, TypeExprPExt);
496ast_payload_map_stub!(AssignTargetP, AssignTargetPExt);
497ast_payload_map_stub!(AssignIdentP, AssignIdentPExt);
498ast_payload_map_stub!(IdentP, IdentPExt);
499ast_payload_map_stub!(ParameterP, ParameterPExt);
500ast_payload_map_stub!(ArgumentP, ArgumentPExt);
501ast_payload_map_stub!(StmtP, StmtPExt);
502ast_payload_map_stub!(StructFieldP, StructFieldPExt);
503ast_payload_map_stub!(FStringP, FStringPExt);