dry4rust/
node.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 std::collections::HashMap;
7use std::collections::HashSet;
8
9/// Kinds of literals — preserves type but erases value.
10#[derive(Debug, Clone, PartialEq, Eq, Hash)]
11pub enum LiteralKind {
12    Int,
13    Float,
14    Str,
15    ByteStr,
16    CStr,
17    Byte,
18    Char,
19    Bool,
20    Null,
21}
22
23/// Kinds of placeholders — what the original identifier referred to.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
25pub enum PlaceholderKind {
26    Variable,
27    Function,
28    Type,
29    Lifetime,
30    Label,
31}
32
33/// Binary operators.
34#[derive(Debug, Clone, PartialEq, Eq, Hash)]
35pub enum BinOpKind {
36    Add,
37    Sub,
38    Mul,
39    Div,
40    Rem,
41    And,
42    Or,
43    BitXor,
44    BitAnd,
45    BitOr,
46    Shl,
47    Shr,
48    Eq,
49    Lt,
50    Le,
51    Ne,
52    Ge,
53    Gt,
54    AddAssign,
55    SubAssign,
56    MulAssign,
57    DivAssign,
58    RemAssign,
59    BitXorAssign,
60    BitAndAssign,
61    BitOrAssign,
62    ShlAssign,
63    ShrAssign,
64    FloorDiv,
65    Pow,
66    In,
67    NotIn,
68    Is,
69    IsNot,
70    FloorDivAssign,
71    PowAssign,
72    Other,
73}
74
75/// Unary operators.
76#[derive(Debug, Clone, PartialEq, Eq, Hash)]
77pub enum UnOpKind {
78    Deref,
79    Not,
80    Neg,
81    Other,
82}
83
84/// The kind of a normalized AST node. Carries only non-child data
85/// (operator kinds, literal kinds, placeholder indices, mutability flags, macro names).
86#[derive(Debug, Clone, PartialEq, Eq, Hash)]
87pub enum NodeKind {
88    // Blocks and statements
89    Block,
90    LetBinding,
91    Semi,
92    Paren,
93
94    // Literals and identifiers
95    Literal(LiteralKind),
96    Placeholder(PlaceholderKind, usize),
97
98    // Operations
99    BinaryOp(BinOpKind),
100    UnaryOp(UnOpKind),
101    Range,
102
103    // Calls and access
104    Call,
105    MethodCall,
106    FieldAccess,
107    Index,
108    Path,
109
110    // Closures and functions
111    Closure,
112    FnSignature,
113
114    // Control flow
115    Return,
116    Break,
117    Continue,
118    Assign,
119
120    // References and pointers
121    Reference {
122        mutable: bool,
123    },
124
125    // Compound types
126    Tuple,
127    Array,
128    Set,
129    Repeat,
130
131    // Type operations
132    Cast,
133    StructInit,
134
135    // Async/error
136    Await,
137    Yield,
138    Try,
139
140    // Control flow structures
141    If,
142    Match,
143    MatchArm,
144    Loop,
145    While,
146    ForLoop,
147    LetExpr,
148
149    // Patterns
150    PatWild,
151    PatPlaceholder(PlaceholderKind, usize),
152    PatTuple,
153    PatStruct,
154    PatOr,
155    PatLiteral,
156    PatReference {
157        mutable: bool,
158    },
159    PatSlice,
160    PatRest,
161    PatRange,
162
163    // Types
164    TypePlaceholder(PlaceholderKind, usize),
165    TypeReference {
166        mutable: bool,
167    },
168    TypeTuple,
169    TypeSlice,
170    TypeArray,
171    TypePath,
172    TypeImplTrait,
173    TypeInfer,
174    TypeUnit,
175    TypeNever,
176
177    // Field initializer (name = value)
178    FieldValue,
179
180    // Macro invocations
181    MacroCall {
182        name: String,
183    },
184
185    // Opaque — unsupported constructs
186    Opaque,
187
188    /// Sentinel for absent optional children, ensuring fixed child positions
189    /// for correct zip alignment in similarity comparison.
190    None,
191}
192
193/// A normalized AST node. Uses a data-driven `{ kind, children }` representation
194/// instead of a large enum with differently-shaped variants. This allows generic
195/// traversal algorithms (count_nodes, reindex, count_matching, extract) to work
196/// without exhaustive matching on every variant.
197///
198/// ## Child ordering conventions
199///
200/// - **Fixed with None sentinels** (always same child count):
201///   - `If` -> [condition, then_branch, else_or_None]
202///   - `LetBinding` -> [pattern, type_or_None, init_or_None, diverge_or_None]
203///   - `Range` / `PatRange` -> [from_or_None, to_or_None]
204///   - `MatchArm` -> [pattern, guard_or_None, body]
205/// - **Fixed children first, variable after** (for zip alignment):
206///   - `Call` -> [func, arg0, arg1, ...]
207///   - `MethodCall` -> [receiver, method, arg0, ...]
208///   - `Closure` -> [body, param0, ...]
209///   - `FnSignature` -> [return_type_or_None, param0, ...]
210///   - `Match` -> [expr, arm0, arm1, ...]
211///   - `StructInit` -> [rest_or_None, field0, field1, ...]
212///   - `MacroCall` -> [arg0, arg1, ...]
213/// - **Variable-length (0 or 1)**: `Return`, `Break` -> `[]` or `[value]`
214/// - **Homogeneous**: `Block`, `Tuple`, `Array`, `Path`, `PatTuple`, etc. -> [elem0, ...]
215/// - **All other fixed**: e.g. `BinaryOp` -> [left, right], `ForLoop` -> [pat, iter, body]
216#[derive(Debug, Clone, PartialEq, Eq, Hash)]
217pub struct NormalizedNode {
218    pub kind: NodeKind,
219    pub children: Vec<Self>,
220}
221
222impl NormalizedNode {
223    /// Create a leaf node (no children).
224    #[must_use]
225    pub const fn leaf(kind: NodeKind) -> Self {
226        Self {
227            kind,
228            children: vec![],
229        }
230    }
231
232    /// Create a node with children.
233    #[must_use]
234    pub const fn with_children(kind: NodeKind, children: Vec<Self>) -> Self {
235        Self { kind, children }
236    }
237
238    /// Create a None sentinel node.
239    #[must_use]
240    pub const fn none() -> Self {
241        Self::leaf(NodeKind::None)
242    }
243
244    /// Convert an `Option<NormalizedNode>` to a node, using the [`Self::none`] sentinel for
245    /// absent values.
246    pub fn opt(node: Option<Self>) -> Self {
247        node.unwrap_or_else(Self::none)
248    }
249
250    /// Check if this is a None sentinel node.
251    #[must_use]
252    pub const fn is_none(&self) -> bool {
253        matches!(self.kind, NodeKind::None)
254    }
255}
256
257// -- Placeholder re-indexing --------------------------------------------------
258
259/// Collects all placeholder occurrences in depth-first order, building
260/// a mapping from (kind, old_index) -> new_sequential_index.
261fn collect_placeholder_order(
262    node: &NormalizedNode,
263    order: &mut Vec<(PlaceholderKind, usize)>,
264    seen: &mut HashSet<(PlaceholderKind, usize)>,
265) {
266    // One `if let` rather than a match whose only arm is an `if`. clippy's
267    // collapsible_match asks for the insert to become a match guard, which
268    // would run a mutation during pattern matching; a let chain keeps the
269    // side effect in the condition where it reads as one.
270    if let NodeKind::Placeholder(kind, idx)
271    | NodeKind::PatPlaceholder(kind, idx)
272    | NodeKind::TypePlaceholder(kind, idx) = &node.kind
273        && seen.insert((*kind, *idx))
274    {
275        order.push((*kind, *idx));
276    }
277    for child in &node.children {
278        collect_placeholder_order(child, order, seen);
279    }
280}
281
282/// Applies the reindex mapping to a node, returning a new node with remapped indices.
283fn apply_reindex(
284    node: &NormalizedNode,
285    mapping: &HashMap<(PlaceholderKind, usize), usize>,
286) -> NormalizedNode {
287    let kind = match &node.kind {
288        NodeKind::Placeholder(kind, idx) => {
289            let new_idx = mapping.get(&(*kind, *idx)).copied().unwrap_or(*idx);
290            NodeKind::Placeholder(*kind, new_idx)
291        }
292        NodeKind::PatPlaceholder(kind, idx) => {
293            let new_idx = mapping.get(&(*kind, *idx)).copied().unwrap_or(*idx);
294            NodeKind::PatPlaceholder(*kind, new_idx)
295        }
296        NodeKind::TypePlaceholder(kind, idx) => {
297            let new_idx = mapping.get(&(*kind, *idx)).copied().unwrap_or(*idx);
298            NodeKind::TypePlaceholder(*kind, new_idx)
299        }
300        other => other.clone(),
301    };
302    let children = node
303        .children
304        .iter()
305        .map(|c| apply_reindex(c, mapping))
306        .collect();
307    NormalizedNode { kind, children }
308}
309
310/// Re-index all placeholders in a sub-tree so that indices start from 0
311/// per kind, assigned by first-occurrence depth-first order.
312/// This allows comparing sub-trees extracted from different function contexts.
313#[must_use]
314pub fn reindex_placeholders(node: &NormalizedNode) -> NormalizedNode {
315    let mut order = Vec::new();
316    let mut seen = HashSet::new();
317    collect_placeholder_order(node, &mut order, &mut seen);
318
319    // Build mapping: (kind, old_index) -> new sequential index per kind
320    let mut counters: HashMap<PlaceholderKind, usize> = HashMap::new();
321    let mut mapping: HashMap<(PlaceholderKind, usize), usize> = HashMap::new();
322    for (kind, old_idx) in order {
323        let counter = counters.entry(kind).or_insert(0);
324        mapping.insert((kind, old_idx), *counter);
325        *counter += 1;
326    }
327
328    apply_reindex(node, &mapping)
329}
330
331/// Count the number of nodes in a normalized tree.
332/// None sentinel nodes are not counted.
333pub fn count_nodes(node: &NormalizedNode) -> usize {
334    if node.is_none() {
335        return 0;
336    }
337    1 + node.children.iter().map(count_nodes).sum::<usize>()
338}
dry4rust/node.rs

dry4rust/
node.rs
