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    match &node.kind {
267        NodeKind::Placeholder(kind, idx)
268        | NodeKind::PatPlaceholder(kind, idx)
269        | NodeKind::TypePlaceholder(kind, idx) => {
270            if seen.insert((*kind, *idx)) {
271                order.push((*kind, *idx));
272            }
273        }
274        _ => {}
275    }
276    for child in &node.children {
277        collect_placeholder_order(child, order, seen);
278    }
279}
280
281/// Applies the reindex mapping to a node, returning a new node with remapped indices.
282fn apply_reindex(
283    node: &NormalizedNode,
284    mapping: &HashMap<(PlaceholderKind, usize), usize>,
285) -> NormalizedNode {
286    let kind = match &node.kind {
287        NodeKind::Placeholder(kind, idx) => {
288            let new_idx = mapping.get(&(*kind, *idx)).copied().unwrap_or(*idx);
289            NodeKind::Placeholder(*kind, new_idx)
290        }
291        NodeKind::PatPlaceholder(kind, idx) => {
292            let new_idx = mapping.get(&(*kind, *idx)).copied().unwrap_or(*idx);
293            NodeKind::PatPlaceholder(*kind, new_idx)
294        }
295        NodeKind::TypePlaceholder(kind, idx) => {
296            let new_idx = mapping.get(&(*kind, *idx)).copied().unwrap_or(*idx);
297            NodeKind::TypePlaceholder(*kind, new_idx)
298        }
299        other => other.clone(),
300    };
301    let children = node
302        .children
303        .iter()
304        .map(|c| apply_reindex(c, mapping))
305        .collect();
306    NormalizedNode { kind, children }
307}
308
309/// Re-index all placeholders in a sub-tree so that indices start from 0
310/// per kind, assigned by first-occurrence depth-first order.
311/// This allows comparing sub-trees extracted from different function contexts.
312#[must_use]
313pub fn reindex_placeholders(node: &NormalizedNode) -> NormalizedNode {
314    let mut order = Vec::new();
315    let mut seen = HashSet::new();
316    collect_placeholder_order(node, &mut order, &mut seen);
317
318    // Build mapping: (kind, old_index) -> new sequential index per kind
319    let mut counters: HashMap<PlaceholderKind, usize> = HashMap::new();
320    let mut mapping: HashMap<(PlaceholderKind, usize), usize> = HashMap::new();
321    for (kind, old_idx) in order {
322        let counter = counters.entry(kind).or_insert(0);
323        mapping.insert((kind, old_idx), *counter);
324        *counter += 1;
325    }
326
327    apply_reindex(node, &mapping)
328}
329
330/// Count the number of nodes in a normalized tree.
331/// None sentinel nodes are not counted.
332pub fn count_nodes(node: &NormalizedNode) -> usize {
333    if node.is_none() {
334        return 0;
335    }
336    1 + node.children.iter().map(count_nodes).sum::<usize>()
337}
dry4rust/node.rs

dry4rust/
node.rs
