rust_split/

reassemble.rs

//! Tier 2 step 2 — reassemble a planned split into output module files.
//!
//! Target: a **binary/crate-root** file (e.g. `src/main.rs`) whose items all
//! live in one module and reference each other by bare name. Splitting them
//! into sibling modules breaks those bare references, so the scheme is:
//!
//! - the crate root keeps the file preamble + imports, the source's own
//!   re-exports (`pub use ...` — dropping or demoting one would narrow the
//!   crate's public API) and content-less `mod name;` declarations verbatim,
//!   declares each part module, re-exports each part with a single glob whose
//!   visibility matches the part's widest item (`pub use` when the part has a
//!   public item, else `pub(crate) use`, and nothing for an `impl`-only part),
//!   keeps local types with inherent impls verbatim (their private fields stay
//!   accessible to children), and keeps `fn main` (the entry point stays here);
//! - each part file copies the imports it references by name (preamble stays at
//!   the root), does `use crate::*` to see its siblings via the root re-exports,
//!   and carries its items with `pub(crate)` visibility bumped on so the
//!   re-exports can see them;
//! - unconditional body-carrying modules are extracted whole, with attributes
//!   and visibility on their declarations; large bodies are split recursively;
//! - conditional modules keep their braced boundaries inline, and oversized
//!   leaves stay intact in their original module. An unmet budget never
//!   justifies dropping a declaration.
//!
//! Residual import/visibility that this mechanical scheme misses is left for the
//! compiler to enumerate (the documented O(n) finish). Unmet LOC budgets are
//! reported alongside the preserved output.

use crate::{Exploded, SplitPlan, plan_split};

mod body_paths;
mod imports;
mod inline_mod;
mod visibility;

use body_paths::*;
use imports::*;
use inline_mod::*;
use visibility::*;

/// Per-part overhead beyond the imports: the `use crate::*;` sibling glob plus
/// blank-line separators.
const PART_OVERHEAD: usize = 3;

/// Plan and reassemble a binary crate-root split targeting files `< max_loc`.
/// Indivisible items and conditional modules stay in their original scope;
/// consult `still_oversized` and `max_loc()` for unmet budget constraints.
/// The item budget reserves the full header LOC and the per-part overhead; since a part now copies only the imports it references (a subset of
/// the header), this reservation is conservative — a packed part stays safely
/// under budget.
pub fn split_bin(exploded: &Exploded, max_loc: usize, root_stem: &str) -> SplitOutput {
    let header_loc: usize = exploded
        .manifest
        .rows
        .iter()
        .filter(|r| matches!(r.kind.as_str(), "use" | "preamble" | "extern_crate"))
        .map(|r| r.loc)
        .sum();
    let item_budget = max_loc.saturating_sub(header_loc + PART_OVERHEAD).max(1);
    let plan = plan_split(exploded, item_budget);
    reassemble_bin(exploded, &plan, root_stem, max_loc)
}

/// A file the split would write.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutputFile {
    /// Path relative to the source file's directory (e.g. `main.rs`, `p00.rs`).
    pub path: String,
    pub contents: String,
    pub loc: usize,
}

/// The reassembled output plus the parts of the plan that could not be met.
#[derive(Debug, Clone)]
pub struct SplitOutput {
    pub files: Vec<OutputFile>,
    /// Retained oversized items and output files at or above the LOC budget.
    pub still_oversized: Vec<String>,
}

impl SplitOutput {
    /// The largest output file's LOC (the budget is met iff this is `< max_loc`
    /// and `still_oversized` is empty).
    pub fn max_loc(&self) -> usize {
        self.files.iter().map(|f| f.loc).max().unwrap_or(0)
    }
}

fn loc(text: &str) -> usize {
    if text.is_empty() {
        0
    } else {
        text.lines().count().max(1)
    }
}

/// How the split file sits in the module tree — fixes the sibling path prefix
/// and where sub-module files go. Re-export visibility is per part in both
/// topologies: `pub use` when the part has a public item (preserving the
/// crate's public API), else no wider than the part's items need.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Topology {
    /// A binary or library crate root (`main.rs`, `lib.rs`): siblings reached
    /// via `crate::`, moved private items raised to `pub(crate)` and
    /// re-exported `pub(crate) use`, `fn main` (if any) kept at the root.
    Bin,
    /// A nested library module (`foo/mod.rs`, or a plain `foo.rs` file module):
    /// sub-modules reached via `super::`, moved private items raised only to
    /// `pub(super)`, and each module re-exported with a single glob — `pub use`
    /// to preserve a public API, `pub(crate) use` for crate-visible items, else
    /// a private `use`. A plain `foo.rs` file module resolves `mod bar;` to
    /// `foo/bar.rs`, so its sub-modules are written into a `foo/` subdir.
    Mod,
}

impl Topology {
    /// The visibility a moved private item is raised to so the original module
    /// and the sibling parts still see it. Parts are children of that module,
    /// so `pub(super)` is enough; at a crate root it equals `pub(crate)`, which
    /// the crate-root output has always used.
    fn sibling_vis(self) -> &'static str {
        match self {
            Topology::Bin => "pub(crate) ",
            Topology::Mod => NESTED_SIBLING_VIS,
        }
    }
}

/// Sibling visibility below any module that is not a crate root: generated
/// parts of a `--module` split and the groups of a nested split.
const NESTED_SIBLING_VIS: &str = "pub(super) ";

/// Plan and reassemble a library module targeting files `< max_loc`;
/// sub-modules see each other via `super::*`. Unmet budgets are reported
/// without discarding declarations. Handles both a directory-owning
/// `foo/mod.rs` (sub-modules become siblings) and a plain `foo.rs` file module
/// (sub-modules go in a `foo/` subdir). `root_stem` is the source file's stem.
pub fn split_mod(exploded: &Exploded, max_loc: usize, root_stem: &str) -> SplitOutput {
    let header_loc: usize = exploded
        .manifest
        .rows
        .iter()
        .filter(|r| matches!(r.kind.as_str(), "use" | "preamble" | "extern_crate"))
        .map(|r| r.loc)
        .sum();
    let item_budget = max_loc.saturating_sub(header_loc + PART_OVERHEAD).max(1);
    let plan = plan_split(exploded, item_budget);
    reassemble(exploded, &plan, root_stem, Topology::Mod, max_loc)
}

/// Reassemble a binary crate-root split (back-compat wrapper).
pub fn reassemble_bin(
    exploded: &Exploded,
    plan: &SplitPlan,
    root_stem: &str,
    max_loc: usize,
) -> SplitOutput {
    reassemble(exploded, plan, root_stem, Topology::Bin, max_loc)
}

/// Reassemble a split for the given module topology. `root_stem` is the source
/// file's stem (`"main"` for `main.rs`, `"mod"` for `foo/mod.rs`). `max_loc` is
/// the ceiling, needed to recursively split an oversized test module.
fn reassemble(
    exploded: &Exploded,
    plan: &SplitPlan,
    root_stem: &str,
    topology: Topology,
    max_loc: usize,
) -> SplitOutput {
    let sibling = match topology {
        Topology::Bin => "crate",
        Topology::Mod => "super",
    };
    // A part is a module *below* the file it came from, whichever topology, so
    // the header imports it copies have to be re-anchored one level up.
    let part_rebase = match topology {
        Topology::Bin => Rebase::CrateRoot,
        Topology::Mod => Rebase::Nested,
    };
    let chunk_text = |i: usize| exploded.chunks[i].text.as_str();
    let row = |i: usize| {
        exploded
            .manifest
            .rows
            .iter()
            .find(|r| r.chunk_index == i)
            .unwrap()
    };

    // Where sub-module files go. A directory-owning module file
    // (`main.rs`/`lib.rs`/`foo/mod.rs`) keeps them as siblings; a plain `foo.rs`
    // file module resolves `mod bar;` to `foo/bar.rs`, so its parts go in `foo/`.
    let sub_prefix = match topology {
        Topology::Bin => String::new(),
        Topology::Mod if matches!(root_stem, "mod" | "lib" | "main") => String::new(),
        Topology::Mod => format!("{root_stem}/"),
    };

    // The crate/module preamble (inner `//!` docs / `#![...]` attrs) stays at
    // the root only — a crate-root `#![feature]` is a hard error in a sub-module
    // file, and a crate-root `#![allow]` already covers descendants. The import
    // chunks (`use`/`extern crate`) are copied into each file, but only the ones
    // that file references by name (see `select_imports`). Re-exports
    // (`pub use ...`) are not imports — the planner routes them to
    // `plan.root_items`, kept at the root verbatim.
    let preamble: String = plan
        .header
        .iter()
        .filter(|&&i| row(i).kind == "preamble")
        .map(|&i| chunk_text(i))
        .collect();
    let import_chunks: Vec<usize> = plan
        .header
        .iter()
        .copied()
        .filter(|&i| row(i).kind != "preamble")
        .collect();

    // Source chunks kept verbatim: re-exports, file-mod declarations, and local
    // types whose inherent impls move into children. The parent's private fields
    // remain visible to those children without changing the type's visibility.
    let root_kept: String = plan.root_items.iter().map(|&i| chunk_text(i)).collect();

    // Module names fixed by the source — extracted body-mods and `mod name;`
    // declarations and retained types keep their own names — plus the root's
    // stem and every name the imports and re-exports bind (`mod fmt;` beside
    // `use std::fmt;` is a conflict); a generated part that would collide is
    // renamed with a `_` suffix.
    let mut reserved: std::collections::BTreeSet<String> =
        plan.mods.iter().map(|&i| row(i).name.clone()).collect();
    reserved.insert(root_stem.to_owned());
    for &i in plan
        .root_items
        .iter()
        .chain(&import_chunks)
        .chain(plan.oversized.iter().map(|o| &o.chunk_index))
    {
        if row(i).kind == "mod" {
            reserved.insert(row(i).name.clone());
        } else if let Some(names) = import_provided_names(chunk_text(i)) {
            reserved.extend(names);
        } else if matches!(
            row(i).kind.as_str(),
            "struct" | "enum" | "union" | "type" | "trait" | "trait_alias"
        ) {
            reserved.insert(row(i).name.clone());
        }
    }

    // Items of the original module that a relative import path can start
    // from; copies of such imports in a part are anchored explicitly.
    let local_names = local_path_roots(exploded);

    let mut files = Vec::new();
    let mut mod_decls = String::new();
    let mut reexports = String::new();
    let mut still_oversized = Vec::new();
    // Names referenced by retained types and other root items, `fn main`, and
    // extracted module bodies whose `use super::*` reaches the root's imports.
    let mut root_refs = Refs::default();
    for &i in &plan.root_items {
        root_refs.add(chunk_text(i));
    }

    // Part modules, named after their cluster's dominant item.
    for part in &plan.parts {
        let mut module = part.name.clone();
        while !reserved.insert(module.clone()) {
            module.push('_');
        }
        let mut body = String::new();
        let mut root_main = None;
        // Widest visibility any moved item contributes to a glob re-export.
        let mut export_vis = ExportVis::None;
        // Names this part references, so only the imports it uses are copied in.
        let mut refs = Refs::default();
        for &i in &part.chunk_indices {
            let r = row(i);
            // `fn main` must stay at the crate root, not move into a submodule.
            if r.kind == "fn" && r.name == "main" {
                root_main = Some(i);
                continue;
            }
            export_vis = export_vis.max(item_export_vis(chunk_text(i)));
            refs.add(chunk_text(i));
            body.push_str(&rebase_body_paths(
                &bump_visibility(chunk_text(i), topology.sibling_vis()),
                part_rebase,
            ));
        }
        // If the part was only `fn main`, it produced no module file.
        if body.trim().is_empty() {
            if let Some(i) = root_main {
                // stash main to emit at root via a sentinel part with no module
                files.push(OutputFile {
                    path: String::from("__root_main__"),
                    contents: chunk_text(i).to_owned(),
                    loc: 0,
                });
            }
            continue;
        }
        if let Some(i) = root_main {
            files.push(OutputFile {
                path: String::from("__root_main__"),
                contents: chunk_text(i).to_owned(),
                loc: 0,
            });
        }
        let part_imports = select_imports(
            &import_chunks,
            &exploded.chunks,
            &refs,
            part_rebase,
            &local_names,
        );
        let contents = part_file(&part_imports, sibling, &body);
        mod_decls.push_str(&format!("mod {module};\n"));
        // One glob re-export per module, visibility chosen by the module's widest
        // item: `pub use` when it has a public item (a `pub use *` caps each item
        // at its own visibility, so it covers internal items in the same line);
        // `pub(crate) use` for crate-visible items; a private `use` for items
        // that were private to a nested module; and nothing when the module
        // exposes no nameable item (e.g. a bare `impl` — a glob re-export would
        // only warn "doesn't reexport anything"). This keeps the crate's public
        // surface (`pub use`) reachable and avoids the redundant double
        // re-export.
        if let Some(line) = reexport_line(export_vis, &module, topology) {
            reexports.push_str(&line);
        }
        let l = loc(&contents);
        files.push(OutputFile {
            path: format!("{sub_prefix}{module}.rs"),
            contents,
            loc: l,
        });
    }

    // Unconditional modules are extracted; conditional modules keep their
    // original braced scope, including mutually exclusive same-named bodies.
    for &mod_index in &plan.mods {
        let chunk = chunk_text(mod_index);
        let name = row(mod_index).name.as_str();
        let Some((decl, inner)) = extract_mod(chunk) else {
            // Conditional or unparseable modules remain intact at the root.
            root_refs.add(chunk);
            mod_decls.push_str(chunk);
            continue;
        };
        mod_decls.push_str(&decl);
        // The body's `use super::*;` reaches the root's imports, so the root
        // must keep the imports the body names.
        root_refs.add_source(&inner);
        let inner_loc = loc(&inner);
        if inner_loc < max_loc {
            files.push(OutputFile {
                path: format!("{sub_prefix}{name}.rs"),
                contents: inner,
                loc: inner_loc,
            });
        } else {
            let (nested, nested_still) = split_nested_mod(name, &inner, max_loc);
            files.extend(nested.into_iter().map(|mut f| {
                f.path = format!("{sub_prefix}{}", f.path);
                f
            }));
            still_oversized.extend(nested_still);
        }
    }

    // Keep oversized leaves at their original scope, without visibility or
    // path changes. Retain imports they reference just as for `main`.
    let mut oversized_kept = String::new();
    for over in &plan.oversized {
        let chunk = chunk_text(over.chunk_index);
        oversized_kept.push_str(chunk);
        root_refs.add(chunk);
        still_oversized.push(format!(
            "{} ({} LOC {}, manual extraction)",
            over.name, over.loc, over.kind
        ));
    }

    // Pull the stashed `fn main` (if any) and build the root file.
    let mut root_main = String::new();
    files.retain(|f| {
        if f.path == "__root_main__" {
            root_main = f.contents.clone();
            false
        } else {
            true
        }
    });

    // The root keeps the preamble, then only the imports its own retained code
    // references — retained types and re-exports, `fn main` (if any), plus
    // extracted module bodies whose `use super::*;` resolves against the root.
    // A root with no references needs no imports.
    if !root_main.trim().is_empty() {
        root_refs.add(&root_main);
    }
    let root_imports = if root_refs.is_empty() {
        String::new()
    } else {
        // The root file *is* the original module: its own imports still mean
        // what they meant.
        select_imports(
            &import_chunks,
            &exploded.chunks,
            &root_refs,
            Rebase::Keep,
            &local_names,
        )
    };

    let mut root = String::new();
    root.push_str(&preamble);
    root.push_str(&root_imports);
    root.push_str(&root_kept);
    root.push('\n');
    root.push_str(&mod_decls);
    root.push('\n');
    root.push_str(&reexports);
    root.push_str(&oversized_kept);
    if !root_main.is_empty() {
        root.push('\n');
        root.push_str(&root_main);
    }
    let root_loc = loc(&root);
    files.insert(
        0,
        OutputFile {
            path: format!("{root_stem}.rs"),
            contents: root,
            loc: root_loc,
        },
    );

    // `include!`/`include_str!`/`include_bytes!` resolve their argument against
    // the directory of the file that contains them, so a body written into a
    // subdirectory needs one `../` per level it moved down. The root file did
    // not move (no `/` in its path), and neither did a part that stayed a
    // sibling. Inserting `../` adds no lines, so `loc` is unaffected.
    for file in &mut files {
        let levels = file.path.matches('/').count();
        if levels == 0 || !file.contents.contains("include") {
            continue;
        }
        file.contents = rebase_includes(&file.contents, levels, &file.path);
    }

    for file in &files {
        if file.loc >= max_loc {
            still_oversized.push(format!(
                "{} ({} LOC, contents preserved; exceeds output budget)",
                file.path, file.loc
            ));
        }
    }

    SplitOutput {
        files,
        still_oversized,
    }
}

/// Assemble a part file: its selected imports, the sibling glob so it can reach
/// the rest of the module via the root re-exports, then the moved items.
fn part_file(imports: &str, sibling: &str, body: &str) -> String {
    if imports.trim().is_empty() {
        format!("use {sibling}::*;\n\n{body}")
    } else {
        format!("{imports}\nuse {sibling}::*;\n\n{body}")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{explode, plan_split};

    mod body_paths;
    mod imports;
    mod inline_mod;
    mod split;
    mod visibility;
}