Skip to main content

rust_split/
lib.rs

1use std::collections::BTreeSet;
2use std::fmt;
3use std::ops::Range;
4
5use quote::ToTokens;
6use serde::{Deserialize, Serialize};
7use syn::spanned::Spanned;
8use syn::visit::{self, Visit};
9
10pub mod plan;
11pub mod reassemble;
12pub use plan::{Oversized, Part, SplitPlan, plan_split};
13pub use reassemble::{OutputFile, SplitOutput, Topology, reassemble_bin, split_bin, split_mod};
14
15pub type Result<T> = std::result::Result<T, Error>;
16
17#[derive(Debug)]
18pub enum Error {
19    Parse(syn::Error),
20    NonCharBoundary { offset: usize },
21}
22
23impl fmt::Display for Error {
24    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
25        match self {
26            Self::Parse(error) => write!(formatter, "parse error: {error}"),
27            Self::NonCharBoundary { offset } => {
28                write!(formatter, "span offset {offset} is not a UTF-8 boundary")
29            }
30        }
31    }
32}
33
34impl std::error::Error for Error {}
35
36impl From<syn::Error> for Error {
37    fn from(error: syn::Error) -> Self {
38        Self::Parse(error)
39    }
40}
41
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct Exploded {
44    pub chunks: Vec<Chunk>,
45    pub manifest: Manifest,
46}
47
48#[derive(Debug, Clone, PartialEq, Eq)]
49pub struct Chunk {
50    pub byte_range: ByteRange,
51    pub text: String,
52}
53
54#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
55pub struct ByteRange {
56    pub start: usize,
57    pub end: usize,
58}
59
60impl From<Range<usize>> for ByteRange {
61    fn from(range: Range<usize>) -> Self {
62        Self {
63            start: range.start,
64            end: range.end,
65        }
66    }
67}
68
69#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
70pub struct LineSpan {
71    pub start: usize,
72    pub end: usize,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
76pub struct Manifest {
77    pub rows: Vec<ManifestRow>,
78}
79
80#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
81pub struct ManifestRow {
82    pub idx: usize,
83    pub chunk_index: usize,
84    pub name: String,
85    pub kind: String,
86    pub byte_range: ByteRange,
87    pub span: ByteRange,
88    pub line_span: LineSpan,
89    pub loc: usize,
90    pub adjacency_hint: Vec<String>,
91}
92
93pub fn explode(src: &str) -> Result<Exploded> {
94    let file = syn::parse_file(src)?;
95    let mut items = Vec::with_capacity(file.items.len());
96    for item in &file.items {
97        let range = checked_range(src, item.span().byte_range())?;
98        items.push(ItemInfo {
99            range,
100            name: item_name(item),
101            kind: item_kind(item).to_owned(),
102            adjacency_candidates: item_adjacency_candidates(item),
103        });
104    }
105
106    if items.is_empty() {
107        if src.is_empty() {
108            return Ok(Exploded {
109                chunks: Vec::new(),
110                manifest: Manifest { rows: Vec::new() },
111            });
112        }
113        return Ok(single_preamble(src));
114    }
115
116    let sibling_names = items
117        .iter()
118        .map(|item| item.name.clone())
119        .collect::<BTreeSet<_>>();
120
121    let mut chunk_ranges = Vec::with_capacity(items.len() + 1);
122    let first_start = if has_inner_file_attr(&src[..items[0].range.start]) {
123        if items[0].range.start > 0 {
124            chunk_ranges.push(ChunkSeed {
125                kind: ChunkSeedKind::Preamble,
126                range: 0..items[0].range.start,
127            });
128        }
129        items[0].range.start
130    } else {
131        0
132    };
133
134    for index in 0..items.len() {
135        let start = if index == 0 {
136            first_start
137        } else {
138            split_gap(src, items[index - 1].range.end, items[index].range.start)
139        };
140        let end = if index + 1 == items.len() {
141            src.len()
142        } else {
143            split_gap(src, items[index].range.end, items[index + 1].range.start)
144        };
145        chunk_ranges.push(ChunkSeed {
146            kind: ChunkSeedKind::Item(index),
147            range: start..end,
148        });
149    }
150
151    let mut chunks = Vec::with_capacity(chunk_ranges.len());
152    let mut rows = Vec::with_capacity(chunk_ranges.len());
153    for seed in chunk_ranges {
154        let chunk_index = chunks.len();
155        let text = src[seed.range.clone()].to_owned();
156        chunks.push(Chunk {
157            byte_range: seed.range.clone().into(),
158            text,
159        });
160
161        match seed.kind {
162            ChunkSeedKind::Preamble => rows.push(preamble_row(src, chunk_index, seed.range)),
163            ChunkSeedKind::Item(item_index) => {
164                let item = &items[item_index];
165                let adjacency_hint = item
166                    .adjacency_candidates
167                    .iter()
168                    .filter(|name| sibling_names.contains(*name) && *name != &item.name)
169                    .cloned()
170                    .collect::<BTreeSet<_>>()
171                    .into_iter()
172                    .collect::<Vec<_>>();
173                rows.push(ManifestRow {
174                    idx: rows.len(),
175                    chunk_index,
176                    name: item.name.clone(),
177                    kind: item.kind.clone(),
178                    byte_range: seed.range.into(),
179                    span: item.range.clone().into(),
180                    line_span: line_span(src, &item.range),
181                    loc: loc(&chunks[chunk_index].text),
182                    adjacency_hint,
183                });
184            }
185        }
186    }
187
188    Ok(Exploded {
189        chunks,
190        manifest: Manifest { rows },
191    })
192}
193
194pub fn manifest_toml(manifest: &Manifest) -> Result<String> {
195    toml::to_string_pretty(manifest).map_err(|error| {
196        Error::Parse(syn::Error::new(
197            proc_macro2::Span::call_site(),
198            format!("manifest serialization failed: {error}"),
199        ))
200    })
201}
202
203#[derive(Debug)]
204struct ItemInfo {
205    range: Range<usize>,
206    name: String,
207    kind: String,
208    adjacency_candidates: BTreeSet<String>,
209}
210
211#[derive(Debug)]
212struct ChunkSeed {
213    kind: ChunkSeedKind,
214    range: Range<usize>,
215}
216
217#[derive(Debug)]
218enum ChunkSeedKind {
219    Preamble,
220    Item(usize),
221}
222
223fn single_preamble(src: &str) -> Exploded {
224    let chunk = Chunk {
225        byte_range: ByteRange {
226            start: 0,
227            end: src.len(),
228        },
229        text: src.to_owned(),
230    };
231    Exploded {
232        chunks: vec![chunk],
233        manifest: Manifest {
234            rows: vec![preamble_row(src, 0, 0..src.len())],
235        },
236    }
237}
238
239fn preamble_row(src: &str, chunk_index: usize, range: Range<usize>) -> ManifestRow {
240    ManifestRow {
241        idx: chunk_index,
242        chunk_index,
243        name: "__preamble".to_owned(),
244        kind: "preamble".to_owned(),
245        byte_range: range.clone().into(),
246        span: range.clone().into(),
247        line_span: line_span(src, &range),
248        loc: loc(&src[range]),
249        adjacency_hint: Vec::new(),
250    }
251}
252
253fn checked_range(src: &str, range: Range<usize>) -> Result<Range<usize>> {
254    if !src.is_char_boundary(range.start) {
255        return Err(Error::NonCharBoundary {
256            offset: range.start,
257        });
258    }
259    if !src.is_char_boundary(range.end) {
260        return Err(Error::NonCharBoundary { offset: range.end });
261    }
262    Ok(range)
263}
264
265fn has_inner_file_attr(gap: &str) -> bool {
266    gap.lines().any(|line| {
267        let line = line.trim_start();
268        line.starts_with("//!") || line.starts_with("#![")
269    })
270}
271
272fn split_gap(src: &str, start: usize, end: usize) -> usize {
273    let gap = &src[start..end];
274    if gap.trim().is_empty() {
275        return end;
276    }
277
278    let mut offset = 0;
279    let mut cut = 0;
280    for line in gap.split_inclusive('\n') {
281        offset += line.len();
282        if line.trim().is_empty() {
283            cut = offset;
284        }
285    }
286    start + cut
287}
288
289fn line_span(src: &str, range: &Range<usize>) -> LineSpan {
290    if src.is_empty() {
291        return LineSpan { start: 1, end: 1 };
292    }
293    let start = byte_to_line(src, range.start);
294    let end_offset = if range.end > range.start {
295        range.end.saturating_sub(1)
296    } else {
297        range.end
298    };
299    let end = byte_to_line(src, end_offset);
300    LineSpan { start, end }
301}
302
303fn byte_to_line(src: &str, offset: usize) -> usize {
304    src[..offset.min(src.len())]
305        .bytes()
306        .filter(|byte| *byte == b'\n')
307        .count()
308        + 1
309}
310
311fn loc(text: &str) -> usize {
312    if text.is_empty() {
313        0
314    } else {
315        text.lines().count().max(1)
316    }
317}
318
319fn item_kind(item: &syn::Item) -> &'static str {
320    match item {
321        syn::Item::Const(_) => "const",
322        syn::Item::Enum(_) => "enum",
323        syn::Item::ExternCrate(_) => "extern_crate",
324        syn::Item::Fn(_) => "fn",
325        syn::Item::ForeignMod(_) => "foreign_mod",
326        syn::Item::Impl(_) => "impl",
327        syn::Item::Macro(item) if is_macro_rules(item) => "macro_rules",
328        syn::Item::Macro(_) => "macro",
329        syn::Item::Mod(_) => "mod",
330        syn::Item::Static(_) => "static",
331        syn::Item::Struct(_) => "struct",
332        syn::Item::Trait(_) => "trait",
333        syn::Item::TraitAlias(_) => "trait_alias",
334        syn::Item::Type(_) => "type",
335        syn::Item::Union(_) => "union",
336        syn::Item::Use(_) => "use",
337        syn::Item::Verbatim(_) => "verbatim",
338        _ => "unknown",
339    }
340}
341
342fn item_name(item: &syn::Item) -> String {
343    match item {
344        syn::Item::Const(item) => item.ident.to_string(),
345        syn::Item::Enum(item) => item.ident.to_string(),
346        syn::Item::ExternCrate(item) => item.ident.to_string(),
347        syn::Item::Fn(item) => item.sig.ident.to_string(),
348        syn::Item::ForeignMod(item) => {
349            let name = item.abi.name.as_ref().map(|name| name.value());
350            match name {
351                Some(name) => format!("extern \"{name}\""),
352                None => "extern".to_owned(),
353            }
354        }
355        syn::Item::Impl(item) => impl_name(item),
356        syn::Item::Macro(item) => macro_name(item),
357        syn::Item::Mod(item) => item.ident.to_string(),
358        syn::Item::Static(item) => item.ident.to_string(),
359        syn::Item::Struct(item) => item.ident.to_string(),
360        syn::Item::Trait(item) => item.ident.to_string(),
361        syn::Item::TraitAlias(item) => item.ident.to_string(),
362        syn::Item::Type(item) => item.ident.to_string(),
363        syn::Item::Union(item) => item.ident.to_string(),
364        syn::Item::Use(item) => use_tree_name(&item.tree),
365        syn::Item::Verbatim(tokens) => tokens.to_string(),
366        _ => "unknown".to_owned(),
367    }
368}
369
370fn impl_name(item: &syn::ItemImpl) -> String {
371    let self_ty = normalize_tokens(&item.self_ty.to_token_stream().to_string());
372    if let Some((_, path, _)) = &item.trait_ {
373        let trait_name = normalize_tokens(&path.to_token_stream().to_string());
374        format!("impl {trait_name} for {self_ty}")
375    } else {
376        format!("impl {self_ty}")
377    }
378}
379
380fn macro_name(item: &syn::ItemMacro) -> String {
381    item.ident
382        .as_ref()
383        .map(ToString::to_string)
384        .unwrap_or_else(|| normalize_tokens(&item.mac.path.to_token_stream().to_string()))
385}
386
387fn is_macro_rules(item: &syn::ItemMacro) -> bool {
388    normalize_tokens(&item.mac.path.to_token_stream().to_string()) == "macro_rules"
389}
390
391fn use_tree_name(tree: &syn::UseTree) -> String {
392    match tree {
393        syn::UseTree::Path(path) => format!("{}::{}", path.ident, use_tree_name(&path.tree)),
394        syn::UseTree::Name(name) => name.ident.to_string(),
395        syn::UseTree::Rename(rename) => format!("{} as {}", rename.ident, rename.rename),
396        syn::UseTree::Glob(_) => "*".to_owned(),
397        syn::UseTree::Group(group) => {
398            let names = group
399                .items
400                .iter()
401                .map(use_tree_name)
402                .collect::<Vec<_>>()
403                .join(", ");
404            format!("{{ {names} }}")
405        }
406    }
407}
408
409fn normalize_tokens(value: &str) -> String {
410    value
411        .replace(" :: ", "::")
412        .replace(" < ", "<")
413        .replace(" >", ">")
414        .replace(" ,", ",")
415        .replace("( ", "(")
416        .replace(" )", ")")
417        .replace("[ ", "[")
418        .replace(" ]", "]")
419}
420
421fn item_adjacency_candidates(item: &syn::Item) -> BTreeSet<String> {
422    let mut visitor = AdjacencyVisitor::default();
423    visitor.visit_item(item);
424    visitor.names
425}
426
427/// Every identifier a single top-level item references (typed visitor + raw
428/// macro idents), used by the reassembler to copy only the imports a part
429/// actually names. `None` when the chunk does not parse as an item, signalling
430/// the caller to keep all imports rather than risk dropping a needed one.
431pub(crate) fn referenced_idents(chunk_text: &str) -> Option<BTreeSet<String>> {
432    let item = syn::parse_str::<syn::Item>(chunk_text).ok()?;
433    Some(item_adjacency_candidates(&item))
434}
435
436/// Every identifier a multi-item source fragment references — used for an
437/// extracted module body, whose `use super::*;` reaches the root's imports, so
438/// the root must keep the imports the body names. `None` when the fragment does
439/// not parse, signalling the caller to keep all imports.
440pub(crate) fn referenced_idents_in_source(src: &str) -> Option<BTreeSet<String>> {
441    let file = syn::parse_str::<syn::File>(src).ok()?;
442    let mut names = BTreeSet::new();
443    for item in &file.items {
444        names.extend(item_adjacency_candidates(item));
445    }
446    Some(names)
447}
448
449/// Recursively collect every identifier in a macro's token stream.
450fn collect_macro_idents(tokens: &proc_macro2::TokenStream, names: &mut BTreeSet<String>) {
451    for tree in tokens.clone() {
452        match tree {
453            proc_macro2::TokenTree::Ident(ident) => {
454                names.insert(ident.to_string());
455            }
456            proc_macro2::TokenTree::Group(group) => collect_macro_idents(&group.stream(), names),
457            _ => {}
458        }
459    }
460}
461
462#[derive(Default)]
463struct AdjacencyVisitor {
464    names: BTreeSet<String>,
465}
466
467impl<'ast> Visit<'ast> for AdjacencyVisitor {
468    fn visit_ident(&mut self, ident: &'ast proc_macro2::Ident) {
469        self.names.insert(ident.to_string());
470    }
471
472    fn visit_item_fn(&mut self, item: &'ast syn::ItemFn) {
473        for input in &item.sig.inputs {
474            self.visit_fn_arg(input);
475        }
476        self.visit_block(&item.block);
477    }
478
479    fn visit_item_const(&mut self, item: &'ast syn::ItemConst) {
480        self.visit_type(&item.ty);
481        self.visit_expr(&item.expr);
482    }
483
484    fn visit_item_static(&mut self, item: &'ast syn::ItemStatic) {
485        self.visit_type(&item.ty);
486        self.visit_expr(&item.expr);
487    }
488
489    fn visit_item_type(&mut self, item: &'ast syn::ItemType) {
490        self.visit_type(&item.ty);
491    }
492
493    fn visit_item_struct(&mut self, item: &'ast syn::ItemStruct) {
494        self.visit_fields(&item.fields);
495    }
496
497    fn visit_item_enum(&mut self, item: &'ast syn::ItemEnum) {
498        for variant in &item.variants {
499            self.visit_fields(&variant.fields);
500            if let Some((_, discriminant)) = &variant.discriminant {
501                self.visit_expr(discriminant);
502            }
503        }
504    }
505
506    fn visit_item_union(&mut self, item: &'ast syn::ItemUnion) {
507        self.visit_fields_named(&item.fields);
508    }
509
510    fn visit_item_trait(&mut self, item: &'ast syn::ItemTrait) {
511        for supertrait in &item.supertraits {
512            self.visit_type_param_bound(supertrait);
513        }
514        for trait_item in &item.items {
515            self.visit_trait_item(trait_item);
516        }
517    }
518
519    fn visit_item_trait_alias(&mut self, item: &'ast syn::ItemTraitAlias) {
520        for bound in &item.bounds {
521            self.visit_type_param_bound(bound);
522        }
523    }
524
525    fn visit_item_impl(&mut self, item: &'ast syn::ItemImpl) {
526        self.visit_type(&item.self_ty);
527        if let Some((_, path, _)) = &item.trait_ {
528            self.visit_path(path);
529        }
530        for impl_item in &item.items {
531            self.visit_impl_item(impl_item);
532        }
533    }
534
535    fn visit_item_macro(&mut self, item: &'ast syn::ItemMacro) {
536        self.visit_macro(&item.mac);
537    }
538
539    fn visit_macro(&mut self, mac: &'ast syn::Macro) {
540        // syn does not parse macro bodies, so references inside `json!{ ... }`,
541        // `vec![ foo() ]`, etc. are invisible to the typed visitor. Harvest
542        // every identifier from the raw token stream; the sibling-name filter
543        // discards the noise, keeping only real cross-references.
544        collect_macro_idents(&mac.tokens, &mut self.names);
545        visit::visit_macro(self, mac);
546    }
547
548    fn visit_item_mod(&mut self, item: &'ast syn::ItemMod) {
549        if let Some((_, items)) = &item.content {
550            for item in items {
551                self.visit_item(item);
552            }
553        }
554    }
555
556    fn visit_pat_ident(&mut self, pat: &'ast syn::PatIdent) {
557        self.names.insert(pat.ident.to_string());
558        visit::visit_pat_ident(self, pat);
559    }
560}