dry4rust/
sub_unit_extractor.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::code_unit::CodeUnitKind;
7use crate::extractor::SubUnit;
8use crate::node::NodeKind;
9use crate::node::NormalizedNode;
10use crate::node::count_nodes;
11use crate::node::reindex_placeholders;
12
13/// Walks a normalized tree and collects the compound structures inside it.
14///
15/// Every node kind that carries a body contributes that body as a candidate
16/// sub-unit, and the walk then descends into all children regardless. What
17/// varies between kinds is only *which* child holds the body, so most of them
18/// reduce to a lookup; `If` and `Match` are the two that do not, because one
19/// has an optional second branch and the other has a variable number.
20pub struct SubUnitExtractor {
21    min_node_count: usize,
22}
23
24impl SubUnitExtractor {
25    #[must_use]
26    pub const fn new(min_node_count: usize) -> Self {
27        Self { min_node_count }
28    }
29
30    /// Collect every sub-unit at or below `node` that meets the node floor.
31    ///
32    /// A node's own contributions come before its children's, and children
33    /// are visited in order, so the result is a pre-order walk.
34    #[must_use]
35    pub fn extract(&self, node: &NormalizedNode) -> Vec<SubUnit> {
36        let mut found = self.own_units(node);
37        for child in &node.children {
38            found.extend(self.extract(child));
39        }
40        found
41    }
42
43    fn own_units(&self, node: &NormalizedNode) -> Vec<SubUnit> {
44        match &node.kind {
45            NodeKind::If => self.if_units(node),
46            NodeKind::Match => self.match_units(node),
47            _ => self.bodied_unit(node),
48        }
49    }
50
51    /// `If` is `[condition, then_branch, else_or_None]`. The else child is
52    /// always present in the normalized form, so an absent else has to be
53    /// told apart from a real one by asking the node rather than the slot.
54    fn if_units(&self, node: &NormalizedNode) -> Vec<SubUnit> {
55        let then_branch = node
56            .children
57            .get(1)
58            .and_then(|n| self.made(n, CodeUnitKind::IfBranch, "if-then branch"));
59        let else_branch = node
60            .children
61            .get(2)
62            .filter(|n| !n.is_none())
63            .and_then(|n| self.made(n, CodeUnitKind::IfBranch, "if-else branch"));
64        then_branch.into_iter().chain(else_branch).collect()
65    }
66
67    /// `Match` is `[expr, arm0, arm1, ...]` and each arm is
68    /// `[pattern, guard_or_None, body]`. Arms are numbered from one in the
69    /// description because that is what a reader counting them would say.
70    fn match_units(&self, node: &NormalizedNode) -> Vec<SubUnit> {
71        node.children
72            .iter()
73            .skip(1)
74            .enumerate()
75            .filter_map(|(i, arm)| {
76                let body = arm.children.get(2)?;
77                self.made(
78                    body,
79                    CodeUnitKind::MatchArm,
80                    &format!("match arm {}", i + 1),
81                )
82            })
83            .collect()
84    }
85
86    fn bodied_unit(&self, node: &NormalizedNode) -> Vec<SubUnit> {
87        let Some(target) = Self::body_target(&node.kind) else {
88            return Vec::new();
89        };
90        let (index, kind, description) = target;
91        node.children
92            .get(index)
93            .and_then(|body| self.made(body, kind, description))
94            .into_iter()
95            .collect()
96    }
97
98    /// Which child holds the body, what to call the unit, and how to describe
99    /// it -- for every kind whose body sits at a fixed position.
100    const fn body_target(kind: &NodeKind) -> Option<(usize, CodeUnitKind, &'static str)> {
101        match kind {
102            NodeKind::Loop => Some((0, CodeUnitKind::LoopBody, "loop body")),
103            NodeKind::While => Some((1, CodeUnitKind::LoopBody, "while body")),
104            NodeKind::ForLoop => Some((2, CodeUnitKind::LoopBody, "for body")),
105            NodeKind::Closure => Some((0, CodeUnitKind::Block, "closure body")),
106            _ => None,
107        }
108    }
109
110    /// Re-index to canonical placeholder form and keep the result only if it
111    /// is big enough to be worth comparing against anything else.
112    fn made(
113        &self,
114        node: &NormalizedNode,
115        kind: CodeUnitKind,
116        description: &str,
117    ) -> Option<SubUnit> {
118        let reindexed = reindex_placeholders(node);
119        let node_count = count_nodes(&reindexed);
120        if node_count < self.min_node_count {
121            return None;
122        }
123        Some(SubUnit {
124            kind,
125            node: reindexed,
126            node_count,
127            description: description.to_string(),
128        })
129    }
130}
dry4rust/sub_unit_extractor.rs

dry4rust/
sub_unit_extractor.rs
