use crate::{Exploded, ManifestRow, naming};
use std::collections::{BTreeMap, BTreeSet, HashMap};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SplitPlan {
pub header: Vec<usize>,
pub root_items: Vec<usize>,
pub mods: Vec<usize>,
pub parts: Vec<Part>,
pub oversized: Vec<Oversized>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Part {
pub name: String,
pub chunk_indices: Vec<usize>,
pub loc: usize,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Oversized {
pub chunk_index: usize,
pub name: String,
pub kind: String,
pub loc: usize,
pub recoverable: bool,
}
impl SplitPlan {
pub fn is_complete(&self) -> bool {
self.oversized.is_empty()
}
}
struct Item {
chunk_index: usize,
name: String,
kind: String,
loc: usize,
refs: Vec<String>,
title: Option<String>,
section: usize,
}
impl Item {
fn new(row: &ManifestRow, chunk_text: &str, section: usize) -> Self {
let lead = row.span.start.saturating_sub(row.byte_range.start);
Self {
chunk_index: row.chunk_index,
name: row.name.clone(),
kind: row.kind.clone(),
loc: row.loc,
refs: row.adjacency_hint.clone(),
title: chunk_text.get(..lead).and_then(naming::banner_title),
section,
}
}
}
pub fn plan_split(exploded: &Exploded, max_loc: usize) -> SplitPlan {
let mut header = Vec::new();
let mut root_items = Vec::new();
let mut mods = Vec::new();
let mut oversized = Vec::new();
let mut items: Vec<Item> = Vec::new();
let owned_types = inherent_type_names(exploded);
let lexical_items = lexical_root_items(exploded, &owned_types);
let mut section = 0;
for row in &exploded.manifest.rows {
let chunk_text = exploded.chunks[row.chunk_index].text.as_str();
let lead = row.span.start.saturating_sub(row.byte_range.start);
if chunk_text
.get(..lead)
.and_then(naming::banner_title)
.is_some()
{
section = row.chunk_index + 1;
}
match row.kind.as_str() {
_ if lexical_items.contains(&row.chunk_index) => root_items.push(row.chunk_index),
"preamble" => header.push(row.chunk_index),
"use" | "extern_crate" => {
if is_reexport(chunk_text) {
root_items.push(row.chunk_index);
} else {
header.push(row.chunk_index);
}
}
"mod" => match mod_body(chunk_text) {
Some(true) => mods.push(row.chunk_index),
Some(false) => root_items.push(row.chunk_index),
None => items.push(Item::new(row, chunk_text, section)),
},
_ if row.loc >= max_loc => oversized.push(Oversized {
chunk_index: row.chunk_index,
name: row.name.clone(),
kind: row.kind.clone(),
loc: row.loc,
recoverable: row.kind == "mod",
}),
"struct" | "enum" | "union" if owned_types.contains(&row.name) => {
root_items.push(row.chunk_index);
}
"impl" if has_owner_local_members(chunk_text) => root_items.push(row.chunk_index),
_ => items.push(Item::new(row, chunk_text, section)),
}
}
items.retain(|item| {
if item.loc >= max_loc {
oversized.push(Oversized {
chunk_index: item.chunk_index,
name: item.name.clone(),
kind: item.kind.clone(),
loc: item.loc,
recoverable: item.kind == "mod",
});
false
} else {
true
}
});
oversized.sort_by_key(|over| over.chunk_index);
let parts = if items.is_empty() {
Vec::new()
} else {
cluster(&items, max_loc)
};
SplitPlan {
header,
root_items,
mods,
parts,
oversized,
}
}
fn has_owner_local_members(text: &str) -> bool {
let Ok(syn::Item::Impl(item)) = syn::parse_str(text) else {
return false;
};
item.trait_.is_none()
&& item.items.iter().any(|member| {
let vis = match member {
syn::ImplItem::Fn(item) => &item.vis,
syn::ImplItem::Const(item) => &item.vis,
syn::ImplItem::Type(item) => &item.vis,
_ => return false,
};
matches!(vis, syn::Visibility::Restricted(vis) if vis.path.is_ident("self"))
})
}
fn lexical_root_items(exploded: &Exploded, owners: &BTreeSet<String>) -> BTreeSet<usize> {
let macros: BTreeMap<usize, String> = exploded
.manifest
.rows
.iter()
.filter(|r| r.kind == "macro_rules")
.map(|r| (r.chunk_index, r.name.clone()))
.collect();
if macros.is_empty() {
return BTreeSet::new();
}
let names: BTreeSet<&str> = macros.values().map(String::as_str).collect();
let mut consumers: BTreeMap<String, Vec<usize>> = BTreeMap::new();
let mut dependencies: BTreeMap<usize, Vec<String>> = BTreeMap::new();
let mut pending = BTreeSet::new();
for row in &exploded.manifest.rows {
if matches!(row.kind.as_str(), "preamble" | "use" | "extern_crate") {
continue;
}
let mut refs = crate::referenced_idents(&exploded.chunks[row.chunk_index].text)
.unwrap_or_else(|| names.iter().map(|s| (*s).to_owned()).collect());
if let Some(name) = macros.get(&row.chunk_index) {
refs.insert(name.clone());
}
let refs: Vec<String> = refs
.into_iter()
.filter(|s| names.contains(s.as_str()))
.collect();
if (matches!(row.kind.as_str(), "struct" | "enum" | "union") && owners.contains(&row.name))
|| (row.kind == "impl"
&& has_owner_local_members(&exploded.chunks[row.chunk_index].text))
{
pending.extend(refs.iter().cloned());
}
for name in &refs {
consumers
.entry(name.clone())
.or_default()
.push(row.chunk_index);
}
dependencies.insert(row.chunk_index, refs);
}
let mut seen = BTreeSet::new();
let mut kept = BTreeSet::new();
while let Some(name) = pending.pop_first() {
if !seen.insert(name.clone()) {
continue;
}
for &index in &consumers[&name] {
if kept.insert(index) {
pending.extend(dependencies[&index].iter().cloned());
}
}
}
kept
}
fn inherent_type_names(exploded: &Exploded) -> std::collections::BTreeSet<String> {
exploded
.manifest
.rows
.iter()
.filter(|row| row.kind == "impl")
.filter_map(|row| {
let syn::Item::Impl(item) =
syn::parse_str(&exploded.chunks[row.chunk_index].text).ok()?
else {
return None;
};
if item.trait_.is_some() {
return None;
}
let syn::Type::Path(ty) = item.self_ty.as_ref() else {
return None;
};
if ty.qself.is_some() || ty.path.leading_colon.is_some() {
return None;
}
let segments = &ty.path.segments;
if segments.len() == 1 || (segments.len() == 2 && segments[0].ident == "self") {
Some(segments.last()?.ident.to_string())
} else {
None
}
})
.collect()
}
fn is_reexport(chunk_text: &str) -> bool {
match syn::parse_str::<syn::Item>(chunk_text) {
Ok(syn::Item::Use(item)) => !matches!(item.vis, syn::Visibility::Inherited),
Ok(syn::Item::ExternCrate(item)) => !matches!(item.vis, syn::Visibility::Inherited),
_ => false,
}
}
fn mod_body(chunk_text: &str) -> Option<bool> {
match syn::parse_str::<syn::Item>(chunk_text) {
Ok(syn::Item::Mod(item)) => Some(item.content.is_some()),
_ => None,
}
}
fn cluster(items: &[Item], max_loc: usize) -> Vec<Part> {
let n = items.len();
let mut name_to_item: HashMap<(usize, &str), usize> = HashMap::new();
for (i, item) in items.iter().enumerate() {
name_to_item
.entry((item.section, item.name.as_str()))
.or_insert(i);
}
let mut weights: HashMap<(usize, usize), usize> = HashMap::new();
for (i, item) in items.iter().enumerate() {
for reference in &item.refs {
if let Some(&j) = name_to_item.get(&(item.section, reference.as_str()))
&& i != j
{
*weights.entry((i.min(j), i.max(j))).or_insert(0) += 1;
}
}
}
let mut edges: Vec<((usize, usize), usize)> = weights.into_iter().collect();
edges.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
let mut component: Vec<usize> = (0..n).collect();
for &((i, j), _weight) in &edges {
let ri = find(&mut component, i);
let rj = find(&mut component, j);
if ri != rj {
component[rj] = ri;
}
}
let mut component_loc: HashMap<usize, usize> = HashMap::new();
for (i, item) in items.iter().enumerate() {
let root = find(&mut component, i);
*component_loc.entry(root).or_insert(0) += item.loc;
}
let target = (max_loc / 2).max(1);
let share = |component_root: usize, component_loc: &HashMap<usize, usize>| -> usize {
let total = component_loc[&component_root];
if total <= target {
total
} else {
total.div_ceil(total.div_ceil(target))
}
};
let mut parent: Vec<usize> = (0..n).collect();
let mut group_loc: Vec<usize> = items.iter().map(|item| item.loc).collect();
for ((i, j), _weight) in edges {
let cap = share(find(&mut component, i), &component_loc);
let ri = find(&mut parent, i);
let rj = find(&mut parent, j);
if ri != rj && group_loc[ri] + group_loc[rj] <= cap {
parent[rj] = ri;
group_loc[ri] += group_loc[rj];
}
}
let mut first_chunk: HashMap<usize, usize> = HashMap::new();
for (i, item) in items.iter().enumerate() {
let root = find(&mut parent, i);
let first = first_chunk.entry(root).or_insert(usize::MAX);
*first = (*first).min(item.chunk_index);
}
let mut by_component: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
for i in 0..n {
if find(&mut parent, i) == i {
let component_root = find(&mut component, i);
by_component.entry(component_root).or_default().push(i);
}
}
for (component_root, mut cluster_roots) in by_component {
let cap = share(component_root, &component_loc);
cluster_roots.sort_by_key(|root| first_chunk[root]);
let mut accumulator: Option<usize> = None;
for root in cluster_roots {
match accumulator {
Some(acc) if group_loc[acc] + group_loc[root] <= cap => {
parent[root] = acc;
group_loc[acc] += group_loc[root];
}
_ => accumulator = Some(root),
}
}
}
let mut groups: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
for i in 0..n {
let root = find(&mut parent, i);
groups.entry(root).or_default().push(i);
}
let mut parts: Vec<Part> = groups
.into_values()
.map(|members| {
let mut chunk_indices: Vec<usize> =
members.iter().map(|&i| items[i].chunk_index).collect();
chunk_indices.sort_unstable();
let loc = members.iter().map(|&i| items[i].loc).sum();
Part {
name: cluster_name(&members, items),
chunk_indices,
loc,
}
})
.collect();
parts.sort_by_key(|part| part.chunk_indices.first().copied().unwrap_or(0));
dedupe_names(&mut parts);
parts
}
fn find(parent: &mut [usize], mut x: usize) -> usize {
while parent[x] != x {
parent[x] = parent[parent[x]]; x = parent[x];
}
x
}
fn cluster_name(members: &[usize], items: &[Item]) -> String {
let banner = members
.iter()
.min_by_key(|&&i| (items[i].title.is_none(), items[i].chunk_index))
.and_then(|&i| items[i].title.as_deref());
if let Some(title) = banner {
return naming::module_name(title);
}
let nameable = |kind: &str| {
matches!(
kind,
"fn" | "struct"
| "enum"
| "const"
| "static"
| "trait"
| "trait_alias"
| "type"
| "union"
| "mod"
)
};
let dominant = members
.iter()
.filter(|&&i| nameable(&items[i].kind))
.max_by_key(|&&i| items[i].loc)
.or_else(|| members.iter().max_by_key(|&&i| items[i].loc));
match dominant {
Some(&i) => naming::module_name(&items[i].name),
None => "group".to_owned(),
}
}
fn dedupe_names(parts: &mut [Part]) {
let mut seen: HashMap<String, usize> = HashMap::new();
for part in parts.iter_mut() {
let count = seen.entry(part.name.clone()).or_insert(0);
*count += 1;
if *count > 1 {
part.name = format!("{}_{}", part.name.trim_end_matches('_'), *count);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::explode;
fn fns(defs: &[(&str, &[&str])]) -> String {
let mut src = String::new();
for (name, calls) in defs {
src.push_str(&format!("fn {name}() {{\n"));
for c in *calls {
src.push_str(&format!(" {c}();\n"));
}
src.push_str("}\n\n");
}
src
}
#[test]
fn unrelated_items_become_separate_files_not_merged() {
let src = fns(&[
("a", &[]),
("b", &[]),
("c", &[]),
("d", &[]),
("e", &[]),
("f", &[]),
("g", &[]),
("h", &[]),
("i", &[]),
("j", &[]),
]);
let exploded = explode(&src).unwrap();
let plan = plan_split(&exploded, 10_000);
assert_eq!(plan.parts.len(), 10, "unrelated items must not be merged");
}
#[test]
fn related_items_cluster_together_unrelated_stay_apart() {
let src = fns(&[
("a1", &["a2"]),
("a2", &["a3"]),
("a3", &[]),
("b1", &["b2"]),
("b2", &[]),
("c", &[]),
]);
let exploded = explode(&src).unwrap();
let plan = plan_split(&exploded, 10_000);
assert_eq!(plan.parts.len(), 3, "two clusters + one loner");
let sizes: Vec<usize> = plan.parts.iter().map(|p| p.chunk_indices.len()).collect();
assert!(sizes.contains(&3) && sizes.contains(&2) && sizes.contains(&1));
}
#[test]
fn a_cohesive_group_over_the_ceiling_is_split() {
let src = fns(&[
("a", &["b"]),
("b", &["c"]),
("c", &["d"]),
("d", &["e"]),
("e", &[]),
]);
let exploded = explode(&src).unwrap();
let plan = plan_split(&exploded, 10);
assert!(plan.parts.len() >= 2, "ceiling forces a split");
for part in &plan.parts {
assert!(part.loc < 10, "part over ceiling: {}", part.loc);
}
}
#[test]
fn imports_go_to_the_header_and_clusters_are_named() {
let src = "use std::fmt;\n\nfn alpha() { beta(); }\n\nfn beta() {}\n";
let exploded = explode(src).unwrap();
let plan = plan_split(&exploded, 10_000);
assert_eq!(plan.header.len(), 1);
assert_eq!(plan.parts.len(), 1, "alpha+beta cluster");
assert!(!plan.parts[0].name.is_empty());
assert!(
plan.parts[0]
.name
.chars()
.all(|c| c.is_ascii_lowercase() || c == '_')
);
}
#[test]
fn an_oversized_leaf_is_reported_not_forced() {
let src =
"fn big() {\n let _a = 1;\n let _b = 2;\n let _c = 3;\n let _d = 4;\n}\n";
let exploded = explode(src).unwrap();
let plan = plan_split(&exploded, 5);
assert_eq!(plan.oversized.len(), 1);
assert!(!plan.oversized[0].recoverable);
assert!(plan.parts.is_empty());
}
#[test]
fn reexports_mod_decls_and_body_mods_are_kept_out_of_clustering() {
let src = "use std::fmt;\n\nmod helpers;\n\npub use helpers::Helper;\n\n\
pub fn api() -> i32 {\n 1\n}\n\n\
#[cfg(test)]\nmod tests {\n fn smoke() {}\n}\n";
let exploded = explode(src).unwrap();
let plan = plan_split(&exploded, 10_000);
let text = |indices: &[usize]| -> String {
indices
.iter()
.map(|&i| exploded.chunks[i].text.as_str())
.collect()
};
let header = text(&plan.header);
assert!(header.contains("use std::fmt;"), "plain import -> header");
assert!(
!header.contains("pub use"),
"re-export is not a header import"
);
let root_items = text(&plan.root_items);
assert!(
root_items.contains("mod helpers;"),
"content-less mod decl stays at the root"
);
assert!(
root_items.contains("pub use helpers::Helper;"),
"pub use re-export stays at the root"
);
assert_eq!(plan.mods.len(), 1, "body mod is extracted, never clustered");
assert!(text(&plan.mods).contains("#[cfg(test)]"));
assert_eq!(plan.parts.len(), 1, "only `api` is left to cluster");
assert!(plan.oversized.is_empty());
}
#[test]
fn a_big_connected_component_is_partitioned_not_packed_to_the_ceiling() {
let mut src = String::new();
for i in 0..12 {
let call = if i < 11 {
format!(" chain{}();\n", i + 1)
} else {
String::new()
};
src.push_str(&format!(
"fn chain{i}() {{\n{call} let _a = 1;\n let _b = 2;\n let _c = 3;\n \
let _d = 4;\n let _e = 5;\n let _f = 6;\n let _g = 7;\n}}\n\n"
));
}
let exploded = explode(&src).unwrap();
let plan = plan_split(&exploded, 100);
assert!(
plan.parts.len() >= 3,
"a ~132-LOC component under a 100 ceiling must give >= 3 parts, got {}: {:?}",
plan.parts.len(),
plan.parts.iter().map(|p| p.loc).collect::<Vec<_>>()
);
for part in &plan.parts {
assert!(
part.loc <= 50,
"part must stay near the balanced share (<= half the ceiling), got {}",
part.loc
);
}
}
#[test]
fn satellites_of_a_full_hub_consolidate_instead_of_one_file_each() {
let mut src = String::from("fn hub() {}\n\n");
for i in 0..20 {
src.push_str(&format!(
"fn sat{i}() {{\n hub();\n let _x = {i};\n}}\n\n"
));
}
let exploded = explode(&src).unwrap();
let plan = plan_split(&exploded, 60);
assert!(
plan.parts.len() <= 6,
"satellites must consolidate, got {} parts",
plan.parts.len()
);
for part in &plan.parts {
assert!(part.loc <= 30, "part over the balanced share: {}", part.loc);
}
}
}