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