Skip to main content

dry4rust/rust/normalizer/
expr.rs

1// Copyright (c) 2026 Matjaz Domen Pecan
2// Copyright 2026 Umberto Gotti <umberto.gotti@umbertogotti.dev>
3// Licensed under the MIT License
4// SPDX-License-Identifier: MIT
5
6use crate::node::{NodeKind, NormalizedNode, PlaceholderKind};
7use crate::normalization_context::NormalizationContext;
8
9use super::helpers::{
10    member_to_string, normalize_bin_op, normalize_lit, normalize_macro, normalize_un_op,
11};
12use super::pat::{normalize_pat, normalize_type};
13use syn::Block;
14use syn::Expr;
15use syn::Stmt;
16
17#[allow(clippy::too_many_lines)]
18pub fn normalize_expr(expr: &Expr, ctx: &mut NormalizationContext) -> NormalizedNode {
19    match expr {
20        Expr::Lit(el) => normalize_lit(&el.lit),
21        Expr::Path(ep) => {
22            if ep.path.segments.len() == 1 {
23                let seg = &ep.path.segments[0];
24                let ident = seg.ident.to_string();
25                let kind = if ident.chars().next().is_some_and(char::is_uppercase) {
26                    PlaceholderKind::Type
27                } else {
28                    PlaceholderKind::Variable
29                };
30                let idx = ctx.placeholder(&ident, kind);
31                NormalizedNode::leaf(NodeKind::Placeholder(kind, idx))
32            } else {
33                let segments: Vec<NormalizedNode> = ep
34                    .path
35                    .segments
36                    .iter()
37                    .map(|seg| {
38                        let ident = seg.ident.to_string();
39                        let kind = if ident.chars().next().is_some_and(char::is_uppercase) {
40                            PlaceholderKind::Type
41                        } else {
42                            PlaceholderKind::Variable
43                        };
44                        let idx = ctx.placeholder(&ident, kind);
45                        NormalizedNode::leaf(NodeKind::Placeholder(kind, idx))
46                    })
47                    .collect();
48                NormalizedNode::with_children(NodeKind::Path, segments)
49            }
50        }
51        // BinaryOp -> [left, right]
52        Expr::Binary(eb) => NormalizedNode::with_children(
53            NodeKind::BinaryOp(normalize_bin_op(&eb.op)),
54            vec![
55                normalize_expr(&eb.left, ctx),
56                normalize_expr(&eb.right, ctx),
57            ],
58        ),
59        // UnaryOp -> [operand]
60        Expr::Unary(eu) => NormalizedNode::with_children(
61            NodeKind::UnaryOp(normalize_un_op(&eu.op)),
62            vec![normalize_expr(&eu.expr, ctx)],
63        ),
64        // Call -> [func, arg0, arg1, ...]
65        Expr::Call(ec) => {
66            let mut children = vec![normalize_expr(&ec.func, ctx)];
67            children.extend(ec.args.iter().map(|a| normalize_expr(a, ctx)));
68            NormalizedNode::with_children(NodeKind::Call, children)
69        }
70        // MethodCall -> [receiver, method, arg0, ...]
71        Expr::MethodCall(emc) => {
72            let method_idx = ctx.placeholder(&emc.method.to_string(), PlaceholderKind::Function);
73            let mut children = vec![
74                normalize_expr(&emc.receiver, ctx),
75                NormalizedNode::leaf(NodeKind::Placeholder(PlaceholderKind::Function, method_idx)),
76            ];
77            children.extend(emc.args.iter().map(|a| normalize_expr(a, ctx)));
78            NormalizedNode::with_children(NodeKind::MethodCall, children)
79        }
80        // FieldAccess -> [base, field]
81        Expr::Field(ef) => {
82            let field_idx =
83                ctx.placeholder(&member_to_string(&ef.member), PlaceholderKind::Variable);
84            NormalizedNode::with_children(
85                NodeKind::FieldAccess,
86                vec![
87                    normalize_expr(&ef.base, ctx),
88                    NormalizedNode::leaf(NodeKind::Placeholder(
89                        PlaceholderKind::Variable,
90                        field_idx,
91                    )),
92                ],
93            )
94        }
95        // Index -> [base, index]
96        Expr::Index(ei) => NormalizedNode::with_children(
97            NodeKind::Index,
98            vec![
99                normalize_expr(&ei.expr, ctx),
100                normalize_expr(&ei.index, ctx),
101            ],
102        ),
103        // Closure -> [body, param0, param1, ...]
104        Expr::Closure(ec) => {
105            let mut children = vec![normalize_expr(&ec.body, ctx)];
106            children.extend(ec.inputs.iter().map(|p| normalize_pat(p, ctx)));
107            NormalizedNode::with_children(NodeKind::Closure, children)
108        }
109        // Return -> [] or [value]
110        Expr::Return(er) => {
111            let children: Vec<_> = er
112                .expr
113                .as_ref()
114                .map(|e| vec![normalize_expr(e, ctx)])
115                .unwrap_or_default();
116            NormalizedNode::with_children(NodeKind::Return, children)
117        }
118        // Break -> [] or [value]
119        Expr::Break(eb) => {
120            let children: Vec<_> = eb
121                .expr
122                .as_ref()
123                .map(|e| vec![normalize_expr(e, ctx)])
124                .unwrap_or_default();
125            NormalizedNode::with_children(NodeKind::Break, children)
126        }
127        Expr::Continue(_) => NormalizedNode::leaf(NodeKind::Continue),
128        // Assign -> [left, right]
129        Expr::Assign(ea) => NormalizedNode::with_children(
130            NodeKind::Assign,
131            vec![
132                normalize_expr(&ea.left, ctx),
133                normalize_expr(&ea.right, ctx),
134            ],
135        ),
136        // Reference -> [expr]
137        Expr::Reference(er) => NormalizedNode::with_children(
138            NodeKind::Reference {
139                mutable: er.mutability.is_some(),
140            },
141            vec![normalize_expr(&er.expr, ctx)],
142        ),
143        Expr::Tuple(et) => NormalizedNode::with_children(
144            NodeKind::Tuple,
145            et.elems.iter().map(|e| normalize_expr(e, ctx)).collect(),
146        ),
147        Expr::Array(ea) => NormalizedNode::with_children(
148            NodeKind::Array,
149            ea.elems.iter().map(|e| normalize_expr(e, ctx)).collect(),
150        ),
151        // Repeat -> [elem, len]
152        Expr::Repeat(er) => NormalizedNode::with_children(
153            NodeKind::Repeat,
154            vec![normalize_expr(&er.expr, ctx), normalize_expr(&er.len, ctx)],
155        ),
156        // Cast -> [expr, ty]
157        Expr::Cast(ec) => NormalizedNode::with_children(
158            NodeKind::Cast,
159            vec![normalize_expr(&ec.expr, ctx), normalize_type(&ec.ty, ctx)],
160        ),
161        // StructInit -> [rest_or_None, field0, field1, ...]
162        Expr::Struct(es) => {
163            let mut children = vec![NormalizedNode::opt(
164                es.rest.as_ref().map(|e| normalize_expr(e, ctx)),
165            )];
166            children.extend(es.fields.iter().map(|f| {
167                let field_idx =
168                    ctx.placeholder(&member_to_string(&f.member), PlaceholderKind::Variable);
169                NormalizedNode::with_children(
170                    NodeKind::FieldValue,
171                    vec![
172                        NormalizedNode::leaf(NodeKind::Placeholder(
173                            PlaceholderKind::Variable,
174                            field_idx,
175                        )),
176                        normalize_expr(&f.expr, ctx),
177                    ],
178                )
179            }));
180            NormalizedNode::with_children(NodeKind::StructInit, children)
181        }
182        // Await -> [expr]
183        Expr::Await(ea) => {
184            NormalizedNode::with_children(NodeKind::Await, vec![normalize_expr(&ea.base, ctx)])
185        }
186        // Try -> [expr]
187        Expr::Try(et) => {
188            NormalizedNode::with_children(NodeKind::Try, vec![normalize_expr(&et.expr, ctx)])
189        }
190        // If -> [condition, then_branch, else_or_None]
191        Expr::If(ei) => NormalizedNode::with_children(
192            NodeKind::If,
193            vec![
194                normalize_expr(&ei.cond, ctx),
195                normalize_block(&ei.then_branch, ctx),
196                NormalizedNode::opt(ei.else_branch.as_ref().map(|(_, e)| normalize_expr(e, ctx))),
197            ],
198        ),
199        // Match -> [expr, arm0, arm1, ...]
200        // Each arm is MatchArm -> [pattern, guard_or_None, body]
201        Expr::Match(em) => {
202            let mut children = vec![normalize_expr(&em.expr, ctx)];
203            children.extend(em.arms.iter().map(|arm| {
204                NormalizedNode::with_children(
205                    NodeKind::MatchArm,
206                    vec![
207                        normalize_pat(&arm.pat, ctx),
208                        NormalizedNode::opt(
209                            arm.guard.as_ref().map(|(_, g)| normalize_expr(g, ctx)),
210                        ),
211                        normalize_expr(&arm.body, ctx),
212                    ],
213                )
214            }));
215            NormalizedNode::with_children(NodeKind::Match, children)
216        }
217        // Loop -> [body]
218        Expr::Loop(el) => {
219            NormalizedNode::with_children(NodeKind::Loop, vec![normalize_block(&el.body, ctx)])
220        }
221        // While -> [condition, body]
222        Expr::While(ew) => NormalizedNode::with_children(
223            NodeKind::While,
224            vec![
225                normalize_expr(&ew.cond, ctx),
226                normalize_block(&ew.body, ctx),
227            ],
228        ),
229        // ForLoop -> [pat, iter, body]
230        Expr::ForLoop(ef) => NormalizedNode::with_children(
231            NodeKind::ForLoop,
232            vec![
233                normalize_pat(&ef.pat, ctx),
234                normalize_expr(&ef.expr, ctx),
235                normalize_block(&ef.body, ctx),
236            ],
237        ),
238        Expr::Block(eb) => normalize_block(&eb.block, ctx),
239        // Paren -> [expr]
240        Expr::Paren(ep) => {
241            NormalizedNode::with_children(NodeKind::Paren, vec![normalize_expr(&ep.expr, ctx)])
242        }
243        // Range -> [from_or_None, to_or_None]
244        Expr::Range(er) => NormalizedNode::with_children(
245            NodeKind::Range,
246            vec![
247                NormalizedNode::opt(er.start.as_ref().map(|e| normalize_expr(e, ctx))),
248                NormalizedNode::opt(er.end.as_ref().map(|e| normalize_expr(e, ctx))),
249            ],
250        ),
251        // LetExpr -> [pat, expr]
252        Expr::Let(el) => NormalizedNode::with_children(
253            NodeKind::LetExpr,
254            vec![normalize_pat(&el.pat, ctx), normalize_expr(&el.expr, ctx)],
255        ),
256        Expr::Macro(em) => normalize_macro(&em.mac, ctx),
257        Expr::Group(eg) => normalize_expr(&eg.expr, ctx),
258        Expr::Unsafe(eu) => normalize_block(&eu.block, ctx),
259        Expr::Const(ec) => normalize_block(&ec.block, ctx),
260        _ => NormalizedNode::leaf(NodeKind::Opaque),
261    }
262}
263
264pub fn normalize_stmt(stmt: &Stmt, ctx: &mut NormalizationContext) -> NormalizedNode {
265    match stmt {
266        // LetBinding -> [pattern, type_or_None, init_or_None, diverge_or_None]
267        Stmt::Local(local) => NormalizedNode::with_children(
268            NodeKind::LetBinding,
269            vec![
270                normalize_pat(&local.pat, ctx),
271                NormalizedNode::none(), // type annotations on let bindings are part of the pattern in syn
272                NormalizedNode::opt(
273                    local
274                        .init
275                        .as_ref()
276                        .map(|init| normalize_expr(&init.expr, ctx)),
277                ),
278                NormalizedNode::opt(
279                    local
280                        .init
281                        .as_ref()
282                        .and_then(|init| init.diverge.as_ref())
283                        .map(|(_, expr)| normalize_expr(expr, ctx)),
284                ),
285            ],
286        ),
287        Stmt::Expr(expr, semi) => {
288            let normalized = normalize_expr(expr, ctx);
289            if semi.is_some() {
290                NormalizedNode::with_children(NodeKind::Semi, vec![normalized])
291            } else {
292                normalized
293            }
294        }
295        Stmt::Item(_) => NormalizedNode::leaf(NodeKind::Opaque),
296        Stmt::Macro(sm) => {
297            let normalized = normalize_macro(&sm.mac, ctx);
298            if sm.semi_token.is_some() {
299                NormalizedNode::with_children(NodeKind::Semi, vec![normalized])
300            } else {
301                normalized
302            }
303        }
304    }
305}
306
307pub fn normalize_block(block: &Block, ctx: &mut NormalizationContext) -> NormalizedNode {
308    NormalizedNode::with_children(
309        NodeKind::Block,
310        block.stmts.iter().map(|s| normalize_stmt(s, ctx)).collect(),
311    )
312}